# Path: jaxopt/_src/linear_operator.py
# class FunctionalLinearOperator:
# 
#   def __init__(self, fun, params):
#     self.fun = functools.partial(fun, params)
# 
#   def __call__(self, x):
#     return self.matvec(x)
# 
#   def matvec(self, x):
#     return self.fun(x)
# 
#   def rmatvec(self, x, y):
#     return self.matvec_and_rmatvec(x, y)[1]
# 
#   def matvec_and_rmatvec(self, x, y):
#     matvec_x, vjp = jax.vjp(self.matvec, x)
#     rmatvec_y, = vjp(y)
#     return matvec_x, rmatvec_y
# 
#   def normal_matvec(self, x):
#     """Computes A^T A x from matvec(x) = A x."""
#     matvec_x, vjp = jax.vjp(self.matvec, x)
#     return vjp(matvec_x)[0]
# 
# 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.linear_operator import FunctionalLinearOperator
from jaxopt._src import test_util
import jax
import jax.numpy as jnp
import numpy as onp
# Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Linear Operator tests."""





class LinearOperatorTest(test_util.JaxoptTestCase):

  def test_matvec_and_rmatvec(self):
    rng = onp.random.RandomState(0)
    A = rng.randn(5, 4)
    matvec = lambda A,x: jnp.dot(A, x)
    x = rng.randn(4)
    y = rng.randn(5)

# Path: jaxopt/_src/linear_operator.py
# class FunctionalLinearOperator:
# 
#   def __init__(self, fun, params):
#     self.fun = functools.partial(fun, params)
# 
#   def __call__(self, x):
#     return self.matvec(x)
# 
#   def matvec(self, x):
#     return self.fun(x)
# 
#   def rmatvec(self, x, y):
#     return self.matvec_and_rmatvec(x, y)[1]
# 
#   def matvec_and_rmatvec(self, x, y):
#     matvec_x, vjp = jax.vjp(self.matvec, x)
#     rmatvec_y, = vjp(y)
#     return matvec_x, rmatvec_y
# 
#   def normal_matvec(self, x):
#     """Computes A^T A x from matvec(x) = A x."""
#     matvec_x, vjp = jax.vjp(self.matvec, x)
#     return vjp(matvec_x)[0]
# 
# 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_operator_test.py
from absl.testing import absltest
from jaxopt._src.linear_operator import FunctionalLinearOperator
from jaxopt._src import test_util
import jax
import jax.numpy as jnp
import numpy as onp

# Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Linear Operator tests."""





class LinearOperatorTest(test_util.JaxoptTestCase):

  def test_matvec_and_rmatvec(self):
    rng = onp.random.RandomState(0)
    A = rng.randn(5, 4)
    matvec = lambda A,x: jnp.dot(A, x)
    x = rng.randn(4)
    y = rng.randn(5)

    linop_A = FunctionalLinearOperator(matvec, A)