# Path: jaxopt/_src/linear_solve.py
# def _materialize_array(matvec, shape, dtype=None):
# def _make_ridge_matvec(matvec: Callable, ridge: float = 0.0):
#   def ridge_matvec(v: Any) -> Any:
# def solve_lu(matvec: Callable, b: jnp.ndarray) -> jnp.ndarray:
# def solve_cholesky(matvec: Callable, b: jnp.ndarray) -> jnp.ndarray:
# def solve_cg(matvec: Callable,
#              b: Any,
#              ridge: Optional[float] = None,
#              init: Optional[Any] = None,
#              **kwargs) -> Any:
# def _rmatvec(matvec, x):
# def _normal_matvec(matvec, x):
# def solve_normal_cg(matvec: Callable,
#                     b: Any,
#                     ridge: Optional[float] = None,
#                     init: Optional[Any] = None,
#                     **kwargs) -> Any:
#   def _matvec(x):
# def solve_gmres(matvec: Callable,
#                 b: Any,
#                 ridge: Optional[float] = None,
#                 init: Optional[Any] = None,
#                 tol: float = 1e-5,
#                 **kwargs) -> Any:
# def solve_bicgstab(matvec: Callable,
#                    b: Any,
#                    ridge: Optional[float] = None,
#                    init: Optional[Any] = None,
#                    **kwargs) -> Any:
#     A = _materialize_array(matvec, b.shape, b.dtype)
#     A = _materialize_array(matvec, b.shape, b.dtype)  # 4d array (tensor)
#     A = A.reshape(-1, b.shape[0] * b.shape[1])  # 2d array (matrix)
#     A = _materialize_array(matvec, b.shape)
#     A = _materialize_array(matvec, b.shape)
# 
# Path: jaxopt/linear_solve.py

# 
# Path: jaxopt/_src/test_util.py
# def ridge_solver(X, y, lam):
# def ridge_solver_jac(X, y, lam, eps=1e-4):
# def lasso_skl(X, y, lam, tol=1e-5, fit_intercept=False):
# def lasso_skl_jac(X, y, lam, tol=1e-5, fit_intercept=False, eps=1e-4):
# def enet_skl(X, y, params_prox, tol=1e-5, fit_intercept=False):
# def enet_skl_jac(X, y, params_prox, tol=1e-5, fit_intercept=False, eps=1e-5):
# def multitask_lasso_skl(X, Y, lam, tol=1e-5):
# def logreg_skl(X, y, lam, tol=1e-5, fit_intercept=False,
#                penalty="l2", multiclass=True):
#   def _reshape_coef(coef):
# def logreg_skl_jac(X, y, lam, tol=1e-5, fit_intercept=False, penalty="l2",
#                    multiclass=True, eps=1e-5):
# def multiclass_linear_svm_skl(X, y, lam, tol=1e-5):
# def multiclass_linear_svm_skl_jac(X, y, lam, tol=1e-5, eps=1e-5):
# def lsq_linear_osp(X, y, bounds, tol=1e-10, max_iter=None):
# def lsq_linear_cube_osp(X, y, l, tol=1e-10, max_iter=None):
# def lsq_linear_cube_osp_jac(X, y, l, eps=1e-5, tol=1e-10, max_iter=None):
# def check_states_have_same_types(state1, state2):
# def _canonicalize_dtype(x64_enabled, dtype):
# def canonicalize_dtype(dtype):
# def _dtype(x):
# def default_tolerance():
# def tolerance(dtype, tol=None):
# def device_under_test():
# def _assert_numpy_allclose(a, b, atol=None, rtol=None, err_msg=''):
# def is_sequence(x):
#   def assertArraysEqual(self, x, y, *, check_dtypes=True, err_msg=''):
#   def assertArraysAllClose(self, x, y, *, check_dtypes=True, atol=None,
#                            rtol=None, err_msg=''):
#   def assertDtypesMatch(self, x, y, *, canonicalize_dtypes=True):
#   def assertAllClose(self, x, y, *, check_dtypes=True, atol=None, rtol=None,
#                      canonicalize_dtypes=True, err_msg=''):
#   XX = jnp.dot(X.T, X)
#   I = jnp.eye(X.shape[1])
# class JaxoptTestCase(parameterized.TestCase):


from absl.testing import absltest
from jaxopt._src import linear_solve as _linear_solve
from jaxopt import linear_solve
from jaxopt._src import test_util
import jax
import jax.numpy as jnp
import numpy as onp
    # Tensor case.
    A = rng.randn(5, 3, 5, 3)
    A_mat = A.reshape(15, 15)
    matvec = lambda x: jnp.dot(A_mat, x.ravel()).reshape(5, 3)
    A2 = _linear_solve._materialize_array(matvec, (5, 3))
    self.assertArraysAllClose(A, A2, atol=1e-3)

  def test_rmatvec(self):
    rng = onp.random.RandomState(0)
    A = rng.randn(5, 5)
    matvec = lambda x: jnp.dot(A, x)
    x = rng.randn(5)
    self.assertArraysAllClose(_linear_solve._rmatvec(matvec, x),
                              jnp.dot(A.T, x))

  def test_normal_matvec(self):
    rng = onp.random.RandomState(0)
    A = rng.randn(5, 5)
    matvec = lambda x: jnp.dot(A, x)
    x = rng.randn(5)
    self.assertArraysAllClose(_linear_solve._normal_matvec(matvec, x),
                              jnp.dot(A.T, jnp.dot(A, x)))

  def test_solve_dense(self):
    rng = onp.random.RandomState(0)

    # Matrix case.
    A = rng.randn(5, 5)
    b = rng.randn(5)
    matvec = lambda x: jnp.dot(A, x)

# Path: jaxopt/_src/linear_solve.py
# def _materialize_array(matvec, shape, dtype=None):
# def _make_ridge_matvec(matvec: Callable, ridge: float = 0.0):
#   def ridge_matvec(v: Any) -> Any:
# def solve_lu(matvec: Callable, b: jnp.ndarray) -> jnp.ndarray:
# def solve_cholesky(matvec: Callable, b: jnp.ndarray) -> jnp.ndarray:
# def solve_cg(matvec: Callable,
#              b: Any,
#              ridge: Optional[float] = None,
#              init: Optional[Any] = None,
#              **kwargs) -> Any:
# def _rmatvec(matvec, x):
# def _normal_matvec(matvec, x):
# def solve_normal_cg(matvec: Callable,
#                     b: Any,
#                     ridge: Optional[float] = None,
#                     init: Optional[Any] = None,
#                     **kwargs) -> Any:
#   def _matvec(x):
# def solve_gmres(matvec: Callable,
#                 b: Any,
#                 ridge: Optional[float] = None,
#                 init: Optional[Any] = None,
#                 tol: float = 1e-5,
#                 **kwargs) -> Any:
# def solve_bicgstab(matvec: Callable,
#                    b: Any,
#                    ridge: Optional[float] = None,
#                    init: Optional[Any] = None,
#                    **kwargs) -> Any:
#     A = _materialize_array(matvec, b.shape, b.dtype)
#     A = _materialize_array(matvec, b.shape, b.dtype)  # 4d array (tensor)
#     A = A.reshape(-1, b.shape[0] * b.shape[1])  # 2d array (matrix)
#     A = _materialize_array(matvec, b.shape)
#     A = _materialize_array(matvec, b.shape)
# 
# Path: jaxopt/linear_solve.py

# 
# Path: jaxopt/_src/test_util.py
# def ridge_solver(X, y, lam):
# def ridge_solver_jac(X, y, lam, eps=1e-4):
# def lasso_skl(X, y, lam, tol=1e-5, fit_intercept=False):
# def lasso_skl_jac(X, y, lam, tol=1e-5, fit_intercept=False, eps=1e-4):
# def enet_skl(X, y, params_prox, tol=1e-5, fit_intercept=False):
# def enet_skl_jac(X, y, params_prox, tol=1e-5, fit_intercept=False, eps=1e-5):
# def multitask_lasso_skl(X, Y, lam, tol=1e-5):
# def logreg_skl(X, y, lam, tol=1e-5, fit_intercept=False,
#                penalty="l2", multiclass=True):
#   def _reshape_coef(coef):
# def logreg_skl_jac(X, y, lam, tol=1e-5, fit_intercept=False, penalty="l2",
#                    multiclass=True, eps=1e-5):
# def multiclass_linear_svm_skl(X, y, lam, tol=1e-5):
# def multiclass_linear_svm_skl_jac(X, y, lam, tol=1e-5, eps=1e-5):
# def lsq_linear_osp(X, y, bounds, tol=1e-10, max_iter=None):
# def lsq_linear_cube_osp(X, y, l, tol=1e-10, max_iter=None):
# def lsq_linear_cube_osp_jac(X, y, l, eps=1e-5, tol=1e-10, max_iter=None):
# def check_states_have_same_types(state1, state2):
# def _canonicalize_dtype(x64_enabled, dtype):
# def canonicalize_dtype(dtype):
# def _dtype(x):
# def default_tolerance():
# def tolerance(dtype, tol=None):
# def device_under_test():
# def _assert_numpy_allclose(a, b, atol=None, rtol=None, err_msg=''):
# def is_sequence(x):
#   def assertArraysEqual(self, x, y, *, check_dtypes=True, err_msg=''):
#   def assertArraysAllClose(self, x, y, *, check_dtypes=True, atol=None,
#                            rtol=None, err_msg=''):
#   def assertDtypesMatch(self, x, y, *, canonicalize_dtypes=True):
#   def assertAllClose(self, x, y, *, check_dtypes=True, atol=None, rtol=None,
#                      canonicalize_dtypes=True, err_msg=''):
#   XX = jnp.dot(X.T, X)
#   I = jnp.eye(X.shape[1])
# class JaxoptTestCase(parameterized.TestCase):


# Path: tests/linear_solve_test.py
from absl.testing import absltest
from jaxopt._src import linear_solve as _linear_solve
from jaxopt import linear_solve
from jaxopt._src import test_util
import jax
import jax.numpy as jnp
import numpy as onp

    # Tensor case.
    A = rng.randn(5, 3, 5, 3)
    A_mat = A.reshape(15, 15)
    matvec = lambda x: jnp.dot(A_mat, x.ravel()).reshape(5, 3)
    A2 = _linear_solve._materialize_array(matvec, (5, 3))
    self.assertArraysAllClose(A, A2, atol=1e-3)

  def test_rmatvec(self):
    rng = onp.random.RandomState(0)
    A = rng.randn(5, 5)
    matvec = lambda x: jnp.dot(A, x)
    x = rng.randn(5)
    self.assertArraysAllClose(_linear_solve._rmatvec(matvec, x),
                              jnp.dot(A.T, x))

  def test_normal_matvec(self):
    rng = onp.random.RandomState(0)
    A = rng.randn(5, 5)
    matvec = lambda x: jnp.dot(A, x)
    x = rng.randn(5)
    self.assertArraysAllClose(_linear_solve._normal_matvec(matvec, x),
                              jnp.dot(A.T, jnp.dot(A, x)))

  def test_solve_dense(self):
    rng = onp.random.RandomState(0)

    # Matrix case.
    A = rng.randn(5, 5)
    b = rng.randn(5)
    matvec = lambda x: jnp.dot(A, x)

    x = linear_solve.solve_lu(matvec, b)