From 05499ae4f025e0e103f8fd9709b047ff2b958c10 Mon Sep 17 00:00:00 2001
From: Neil Girdhar <mistersheik@gmail.com>
Date: Thu, 4 Jan 2024 12:35:15 +0700
Subject: [PATCH] Remove unused imports

---
 benchmarks/lbfgs_benchmark.py                  | 2 --
 benchmarks/proximal_gradient_benchmark.py      | 1 -
 examples/deep_learning/plot_sgd_solvers.py     | 2 --
 examples/fixed_point/deep_equilibrium_model.py | 6 +-----
 examples/fixed_point/plot_picard_ode.py        | 3 ---
 examples/implicit_diff/lasso_implicit_diff.py  | 1 -
 jaxopt/_src/anderson.py                        | 9 +++------
 jaxopt/_src/anderson_wrapper.py                | 4 +---
 jaxopt/_src/backtracking_linesearch.py         | 1 -
 jaxopt/_src/base.py                            | 1 -
 jaxopt/_src/bisection.py                       | 2 --
 jaxopt/_src/block_cd.py                        | 4 ----
 jaxopt/_src/cd_qp.py                           | 1 -
 jaxopt/_src/cvxpy_wrapper.py                   | 4 ----
 jaxopt/_src/fixed_point_iteration.py           | 1 -
 jaxopt/_src/gradient_descent.py                | 2 --
 jaxopt/_src/iterative_refinement.py            | 7 +------
 jaxopt/_src/lbfgs.py                           | 1 -
 jaxopt/_src/levenberg_marquardt.py             | 3 +--
 jaxopt/_src/linear_operator.py                 | 3 +--
 jaxopt/_src/nonlinear_cg.py                    | 1 -
 jaxopt/_src/osqp.py                            | 3 +--
 jaxopt/_src/scipy_wrappers.py                  | 2 --
 jaxopt/_src/test_util.py                       | 3 ---
 jaxopt/_src/zoom_linesearch.py                 | 1 -
 tests/anderson_test.py                         | 6 +-----
 tests/anderson_wrapper_test.py                 | 9 ---------
 tests/base_test.py                             | 7 -------
 tests/common_test.py                           | 2 --
 tests/cvxpy_wrapper_test.py                    | 1 -
 tests/eq_qp_test.py                            | 1 -
 tests/fixed_point_iteration_test.py            | 6 +-----
 tests/implicit_diff_test.py                    | 1 -
 tests/isotonic_test.py                         | 1 -
 tests/linear_operator_test.py                  | 1 -
 tests/projected_gradient_test.py               | 1 -
 tests/proximal_gradient_test.py                | 5 -----
 37 files changed, 11 insertions(+), 98 deletions(-)

diff --git a/benchmarks/lbfgs_benchmark.py b/benchmarks/lbfgs_benchmark.py
index 0c1d05a8..f104125e 100644
--- a/benchmarks/lbfgs_benchmark.py
+++ b/benchmarks/lbfgs_benchmark.py
@@ -14,14 +14,12 @@
 
 """Benchmark LBFGS implementation."""
 
-import time
 
 from absl import app
 from absl import flags
 
 from sklearn import datasets
 
-import jax
 import jax.numpy as jnp
 import jaxopt
 
diff --git a/benchmarks/proximal_gradient_benchmark.py b/benchmarks/proximal_gradient_benchmark.py
index 0a9b54a2..7cffb333 100644
--- a/benchmarks/proximal_gradient_benchmark.py
+++ b/benchmarks/proximal_gradient_benchmark.py
@@ -16,7 +16,6 @@
 
 import time
 
-from typing import NamedTuple
 from typing import Sequence
 
 from absl import app
diff --git a/examples/deep_learning/plot_sgd_solvers.py b/examples/deep_learning/plot_sgd_solvers.py
index e9e711b0..679c8c62 100644
--- a/examples/deep_learning/plot_sgd_solvers.py
+++ b/examples/deep_learning/plot_sgd_solvers.py
@@ -38,7 +38,6 @@
 
 from absl import flags
 
-import logging
 import sys
 from timeit import default_timer as timer
 
@@ -49,7 +48,6 @@
 from jaxopt import ArmijoSGD
 from jaxopt import PolyakSGD
 from jaxopt import OptaxSolver
-from jaxopt.tree_util import tree_l2_norm, tree_sub
 
 import optax
 from flax import linen as nn
diff --git a/examples/fixed_point/deep_equilibrium_model.py b/examples/fixed_point/deep_equilibrium_model.py
index 0faf2455..40c8d01d 100644
--- a/examples/fixed_point/deep_equilibrium_model.py
+++ b/examples/fixed_point/deep_equilibrium_model.py
@@ -34,19 +34,16 @@
 """
 
 from functools import partial
-from typing import Any, Mapping, Tuple, Callable
+from typing import Any, Tuple, Callable
 
 from absl import app
 from absl import flags
 
-import flax
 from flax import linen as nn
 
 import jax
 import jax.numpy as jnp
-from jax.tree_util import tree_structure
 
-import jaxopt
 from jaxopt import loss
 from jaxopt import OptaxSolver
 from jaxopt import FixedPointIteration
@@ -58,7 +55,6 @@
 
 import tensorflow_datasets as tfds
 import tensorflow as tf
-from collections import namedtuple
 
 
 dataset_names = [
diff --git a/examples/fixed_point/plot_picard_ode.py b/examples/fixed_point/plot_picard_ode.py
index f9289f66..ea2ceb85 100644
--- a/examples/fixed_point/plot_picard_ode.py
+++ b/examples/fixed_point/plot_picard_ode.py
@@ -58,12 +58,9 @@
 
 from jaxopt import AndersonAcceleration
 
-from jaxopt import objective
-from jaxopt.tree_util import tree_scalar_mul, tree_sub
 
 import numpy as np
 import matplotlib.pyplot as plt
-from sklearn import datasets
 from matplotlib.pyplot import cm
 import scipy.integrate
 
diff --git a/examples/implicit_diff/lasso_implicit_diff.py b/examples/implicit_diff/lasso_implicit_diff.py
index a66c34cd..f6da6164 100644
--- a/examples/implicit_diff/lasso_implicit_diff.py
+++ b/examples/implicit_diff/lasso_implicit_diff.py
@@ -20,7 +20,6 @@
 from absl import app
 from absl import flags
 
-import jax
 import jax.numpy as jnp
 
 from jaxopt import BlockCoordinateDescent
diff --git a/jaxopt/_src/anderson.py b/jaxopt/_src/anderson.py
index 15f0c114..ef153bb2 100644
--- a/jaxopt/_src/anderson.py
+++ b/jaxopt/_src/anderson.py
@@ -17,7 +17,6 @@
 from typing import Any
 from typing import Callable
 from typing import NamedTuple
-from typing import List
 from typing import Union
 
 from typing import Optional
@@ -27,12 +26,10 @@
 import jax.numpy as jnp
 
 from jaxopt._src import base
-from jaxopt._src import linear_solve
-from jaxopt._src.tree_util import tree_l2_norm, tree_sub
-from jaxopt._src.tree_util import tree_vdot, tree_add
-from jaxopt._src.tree_util import tree_mul, tree_scalar_mul
+from jaxopt._src.tree_util import tree_sub
+from jaxopt._src.tree_util import tree_vdot
 from jaxopt._src.tree_util import tree_average, tree_add_scalar_mul
-from jaxopt._src.tree_util import tree_map, tree_gram
+from jaxopt._src.tree_util import tree_map
 
 
 def minimize_residuals(residual_gram, ridge):
diff --git a/jaxopt/_src/anderson_wrapper.py b/jaxopt/_src/anderson_wrapper.py
index dea00a0e..90d22647 100644
--- a/jaxopt/_src/anderson_wrapper.py
+++ b/jaxopt/_src/anderson_wrapper.py
@@ -18,7 +18,6 @@
 from typing import Callable
 from typing import Optional
 from typing import NamedTuple
-from typing import Union
 
 from dataclasses import dataclass
 
@@ -26,8 +25,7 @@
 import jax.numpy as jnp
 
 from jaxopt._src import base
-from jaxopt._src.tree_util import tree_l2_norm, tree_sub, tree_map
-from jaxopt._src.anderson import AndersonAcceleration
+from jaxopt._src.tree_util import tree_sub, tree_map
 from jaxopt._src.anderson import anderson_step, update_history
 
 
diff --git a/jaxopt/_src/backtracking_linesearch.py b/jaxopt/_src/backtracking_linesearch.py
index 72db4722..03ebc7e6 100644
--- a/jaxopt/_src/backtracking_linesearch.py
+++ b/jaxopt/_src/backtracking_linesearch.py
@@ -22,7 +22,6 @@
 
 from dataclasses import dataclass
 
-import jax
 import jax.numpy as jnp
 
 from jaxopt._src import base
diff --git a/jaxopt/_src/base.py b/jaxopt/_src/base.py
index e3b7a507..d5130135 100644
--- a/jaxopt/_src/base.py
+++ b/jaxopt/_src/base.py
@@ -35,7 +35,6 @@
 # jaxopt._src.linear_solve instead.
 # This allows to define linear solver with base.Solver interface,
 # and then exporting them in jaxopt.linear_solve.
-from jaxopt._src import linear_solve
 
 from jaxopt import loop
 from jaxopt import tree_util
diff --git a/jaxopt/_src/bisection.py b/jaxopt/_src/bisection.py
index 943a20d3..97f259ca 100644
--- a/jaxopt/_src/bisection.py
+++ b/jaxopt/_src/bisection.py
@@ -25,8 +25,6 @@
 import jax.numpy as jnp
 
 from jaxopt._src import base
-from jaxopt._src import implicit_diff as idf
-from jaxopt._src import loop
 
 
 class BisectionState(NamedTuple):
diff --git a/jaxopt/_src/block_cd.py b/jaxopt/_src/block_cd.py
index d77eb4d9..7f36eb83 100644
--- a/jaxopt/_src/block_cd.py
+++ b/jaxopt/_src/block_cd.py
@@ -20,7 +20,6 @@
 from typing import Callable
 from typing import NamedTuple
 from typing import Optional
-from typing import Union
 
 from dataclasses import dataclass
 
@@ -28,10 +27,7 @@
 import jax.numpy as jnp
 
 from jaxopt._src import base
-from jaxopt._src import implicit_diff as idf
-from jaxopt._src import loop
 from jaxopt._src import objective
-from jaxopt._src import tree_util
 
 
 class BlockCDState(NamedTuple):
diff --git a/jaxopt/_src/cd_qp.py b/jaxopt/_src/cd_qp.py
index bd994522..e1598796 100644
--- a/jaxopt/_src/cd_qp.py
+++ b/jaxopt/_src/cd_qp.py
@@ -26,7 +26,6 @@
 
 from jaxopt._src import base
 from jaxopt._src import projection
-from jaxopt._src import tree_util
 
 
 class BoxCDQPState(NamedTuple):
diff --git a/jaxopt/_src/cvxpy_wrapper.py b/jaxopt/_src/cvxpy_wrapper.py
index 4a0a246e..cbe99449 100644
--- a/jaxopt/_src/cvxpy_wrapper.py
+++ b/jaxopt/_src/cvxpy_wrapper.py
@@ -13,18 +13,14 @@
 # limitations under the License.
 """CVXPY wrappers."""
 
-from typing import Any
 from typing import Callable
 from typing import Optional
-from typing import Tuple
 
 from dataclasses import dataclass
 
-import jax
 import jax.numpy as jnp 
 from jaxopt._src import base
 from jaxopt._src import implicit_diff as idf
-from jaxopt._src import linear_solve
 from jaxopt._src import tree_util
 
 
diff --git a/jaxopt/_src/fixed_point_iteration.py b/jaxopt/_src/fixed_point_iteration.py
index 24190449..ede40516 100644
--- a/jaxopt/_src/fixed_point_iteration.py
+++ b/jaxopt/_src/fixed_point_iteration.py
@@ -23,7 +23,6 @@
 from dataclasses import dataclass
 
 import jax.numpy as jnp
-from jax.tree_util import tree_leaves, tree_structure
 
 from jaxopt._src import base
 from jaxopt._src.tree_util import tree_l2_norm, tree_sub
diff --git a/jaxopt/_src/gradient_descent.py b/jaxopt/_src/gradient_descent.py
index bbaf100f..e17d4db6 100644
--- a/jaxopt/_src/gradient_descent.py
+++ b/jaxopt/_src/gradient_descent.py
@@ -15,9 +15,7 @@
 """Implementation of gradient descent in JAX."""
 
 from typing import Any
-from typing import Callable
 from typing import NamedTuple
-from typing import Union
 
 from dataclasses import dataclass
 
diff --git a/jaxopt/_src/iterative_refinement.py b/jaxopt/_src/iterative_refinement.py
index 70667aad..9c3ddf19 100644
--- a/jaxopt/_src/iterative_refinement.py
+++ b/jaxopt/_src/iterative_refinement.py
@@ -23,20 +23,15 @@
 from typing import Callable
 from typing import NamedTuple
 from typing import Optional
-from typing import Union
 
 from dataclasses import dataclass
 from functools import partial
 
-import jax
 import jax.numpy as jnp
 
-from jaxopt._src import loop
 from jaxopt._src import base
-from jaxopt._src import implicit_diff as idf
 from jaxopt._src.tree_util import tree_zeros_like, tree_add, tree_sub
-from jaxopt._src.tree_util import tree_add_scalar_mul, tree_scalar_mul
-from jaxopt._src.tree_util import tree_vdot, tree_negative, tree_l2_norm
+from jaxopt._src.tree_util import tree_l2_norm
 from jaxopt._src.linear_operator import _make_linear_operator
 import jaxopt._src.linear_solve as linear_solve
 
diff --git a/jaxopt/_src/lbfgs.py b/jaxopt/_src/lbfgs.py
index 1be37e1b..78c3b54b 100644
--- a/jaxopt/_src/lbfgs.py
+++ b/jaxopt/_src/lbfgs.py
@@ -17,7 +17,6 @@
 import warnings
 
 from dataclasses import dataclass
-from functools import partial
 from typing import Any, Callable, NamedTuple, Optional, Union
 
 import jax
diff --git a/jaxopt/_src/levenberg_marquardt.py b/jaxopt/_src/levenberg_marquardt.py
index 1a3646ac..b02f5520 100644
--- a/jaxopt/_src/levenberg_marquardt.py
+++ b/jaxopt/_src/levenberg_marquardt.py
@@ -14,7 +14,6 @@
 
 """Levenberg-Marquardt algorithm in JAX."""
 
-import math
 from typing import Any
 from typing import Callable
 from typing import Literal
@@ -35,7 +34,7 @@
 from jaxopt._src.linear_solve import solve_inv
 from jaxopt._src.linear_solve import solve_lu
 from jaxopt._src.linear_solve import solve_qr
-from jaxopt._src.tree_util import tree_l2_norm, tree_inf_norm, tree_sub, tree_add, tree_mul
+from jaxopt._src.tree_util import tree_l2_norm, tree_inf_norm
 
 
 class LevenbergMarquardtState(NamedTuple):
diff --git a/jaxopt/_src/linear_operator.py b/jaxopt/_src/linear_operator.py
index a7d20313..72379e20 100644
--- a/jaxopt/_src/linear_operator.py
+++ b/jaxopt/_src/linear_operator.py
@@ -16,9 +16,8 @@
 import functools
 import jax
 import jax.numpy as jnp
-import numpy as onp
 
-from jaxopt.tree_util import tree_map, tree_sum, tree_mul
+from jaxopt.tree_util import tree_map
 
 
 class DenseLinearOperator:
diff --git a/jaxopt/_src/nonlinear_cg.py b/jaxopt/_src/nonlinear_cg.py
index 398f3db2..fa9f412e 100644
--- a/jaxopt/_src/nonlinear_cg.py
+++ b/jaxopt/_src/nonlinear_cg.py
@@ -25,7 +25,6 @@
 from jaxopt._src.linesearch_util import _setup_linesearch
 from jaxopt._src.tree_util import tree_single_dtype, get_real_dtype
 from jaxopt.tree_util import tree_add_scalar_mul
-from jaxopt.tree_util import tree_div
 from jaxopt.tree_util import tree_l2_norm
 from jaxopt.tree_util import tree_scalar_mul
 from jaxopt.tree_util import tree_sub
diff --git a/jaxopt/_src/osqp.py b/jaxopt/_src/osqp.py
index a62c7f63..9d15e5bf 100644
--- a/jaxopt/_src/osqp.py
+++ b/jaxopt/_src/osqp.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 """GPU-friendly implementation of OSQP."""
-from abc import ABC, abstractmethod
+from abc import abstractmethod
 from dataclasses import dataclass
 from functools import partial
 
@@ -24,7 +24,6 @@
 from typing import Union
 
 import jax
-import jax.nn as nn
 import jax.numpy as jnp
 from jax.tree_util import tree_reduce
 
diff --git a/jaxopt/_src/scipy_wrappers.py b/jaxopt/_src/scipy_wrappers.py
index d7c57825..fcd116e5 100644
--- a/jaxopt/_src/scipy_wrappers.py
+++ b/jaxopt/_src/scipy_wrappers.py
@@ -21,8 +21,6 @@
 #   currently only ScipyMinimize exposes this option.
 """
 
-import abc
-import dataclasses
 from dataclasses import dataclass
 from typing import Any
 from typing import Callable
diff --git a/jaxopt/_src/test_util.py b/jaxopt/_src/test_util.py
index 799cfa38..5b41fe33 100644
--- a/jaxopt/_src/test_util.py
+++ b/jaxopt/_src/test_util.py
@@ -14,7 +14,6 @@
 
 """Test utilities."""
 
-from absl.testing import absltest
 from absl.testing import parameterized
 
 import functools
@@ -22,8 +21,6 @@
 import jax
 import jax.numpy as jnp
 
-from jaxopt._src import base
-from jaxopt._src import loss
 
 import numpy as onp
 import scipy as osp
diff --git a/jaxopt/_src/zoom_linesearch.py b/jaxopt/_src/zoom_linesearch.py
index de732e67..96b790c9 100644
--- a/jaxopt/_src/zoom_linesearch.py
+++ b/jaxopt/_src/zoom_linesearch.py
@@ -34,7 +34,6 @@
 from jaxopt.tree_util import tree_scalar_mul
 from jaxopt.tree_util import tree_vdot_real
 from jaxopt.tree_util import tree_conj
-from jaxopt.tree_util import tree_l2_norm
 
 # pylint: disable=g-bare-generic
 # pylint: disable=invalid-name
diff --git a/tests/anderson_test.py b/tests/anderson_test.py
index ec43dade..cf960f2f 100644
--- a/tests/anderson_test.py
+++ b/tests/anderson_test.py
@@ -19,18 +19,14 @@
 import jax
 import jax.numpy as jnp
 from jax import scipy as jsp
-from jax.tree_util import tree_map, tree_all
 from jax.test_util import check_grads
 
 import jaxopt
-from jaxopt.tree_util import tree_l2_norm, tree_scalar_mul
-from jaxopt._src.tree_util import tree_average, tree_sub
-from jaxopt import objective
+from jaxopt.tree_util import tree_l2_norm
 from jaxopt import AndersonAcceleration
 from jaxopt._src import test_util
 
 import numpy as onp
-from sklearn import datasets
 
 
 class AndersonAccelerationTest(test_util.JaxoptTestCase):
diff --git a/tests/anderson_wrapper_test.py b/tests/anderson_wrapper_test.py
index f9a58fa6..19e11b14 100644
--- a/tests/anderson_wrapper_test.py
+++ b/tests/anderson_wrapper_test.py
@@ -14,31 +14,22 @@
 
 
 from absl.testing import absltest
-from absl.testing import parameterized
 
-import jax
 import jax.numpy as jnp
-from jax import config
-from jax.tree_util import tree_map, tree_all
 from jax.test_util import check_grads
 import optax
 
-from jaxopt.tree_util import tree_l2_norm, tree_scalar_mul, tree_sub
 from jaxopt import objective
 
-from jaxopt import projection
 from jaxopt import prox
 from jaxopt._src import test_util
 
 from jaxopt import AndersonWrapper
 from jaxopt import BlockCoordinateDescent
-from jaxopt import GradientDescent
 from jaxopt import OptaxSolver
 from jaxopt import PolyakSGD
 from jaxopt import ProximalGradient
 
-import numpy as onp
-import scipy
 from sklearn import datasets
 
 
diff --git a/tests/base_test.py b/tests/base_test.py
index 42305c71..fca39ba7 100644
--- a/tests/base_test.py
+++ b/tests/base_test.py
@@ -25,7 +25,6 @@
 from typing import Any
 from typing import Callable
 from typing import NamedTuple
-from typing import Optional
 
 import dataclasses
 
@@ -33,12 +32,6 @@
 import jax.numpy as jnp
 
 from jaxopt._src import base
-from jaxopt.tree_util import tree_add
-from jaxopt.tree_util import tree_add_scalar_mul
-from jaxopt.tree_util import tree_l2_norm
-from jaxopt.tree_util import tree_scalar_mul
-from jaxopt.tree_util import tree_sub
-from jaxopt.tree_util import tree_zeros_like
 
 
 class DummySolverState(NamedTuple):
diff --git a/tests/common_test.py b/tests/common_test.py
index 45d71f01..4412e031 100644
--- a/tests/common_test.py
+++ b/tests/common_test.py
@@ -25,7 +25,6 @@
 
 import jax
 import jax.numpy as jnp
-import numpy as onp
 import optax
 from sklearn import datasets
 from sklearn import preprocessing
@@ -39,7 +38,6 @@
 from jaxopt._src.base import LinearOperator
 from jaxopt._src.objective import LeastSquares, MulticlassLinearSvmDual
 from jaxopt._src.osqp import BoxOSQP
-from jaxopt._src.osqp import OSQP
 
 
 N_CALLS = 0
diff --git a/tests/cvxpy_wrapper_test.py b/tests/cvxpy_wrapper_test.py
index 9d325bc9..85df5bfd 100644
--- a/tests/cvxpy_wrapper_test.py
+++ b/tests/cvxpy_wrapper_test.py
@@ -15,7 +15,6 @@
 """CVXPY tests."""
 
 from absl.testing import absltest
-from absl.testing import parameterized
 
 import jax
 import jax.numpy as jnp
diff --git a/tests/eq_qp_test.py b/tests/eq_qp_test.py
index c103e96a..b209c21f 100644
--- a/tests/eq_qp_test.py
+++ b/tests/eq_qp_test.py
@@ -22,7 +22,6 @@
 from jaxopt.base import KKTSolution
 from jaxopt import EqualityConstrainedQP
 from jaxopt._src import test_util
-from jaxopt._src.tree_util import tree_negative
 
 import numpy as onp
 
diff --git a/tests/fixed_point_iteration_test.py b/tests/fixed_point_iteration_test.py
index 630739a0..b6da9012 100644
--- a/tests/fixed_point_iteration_test.py
+++ b/tests/fixed_point_iteration_test.py
@@ -17,16 +17,12 @@
 
 import jax
 import jax.numpy as jnp
-from jax.tree_util import tree_map, tree_all
 from jax.test_util import check_grads
 
-from jaxopt.tree_util import tree_l2_norm, tree_scalar_mul, tree_sub
-from jaxopt import objective
+from jaxopt.tree_util import tree_l2_norm
 from jaxopt import FixedPointIteration
 from jaxopt._src import test_util
 
-import numpy as onp
-from sklearn import datasets
 
 
 class FixedPointIterationTest(test_util.JaxoptTestCase):
diff --git a/tests/implicit_diff_test.py b/tests/implicit_diff_test.py
index ac928b0c..ff204984 100644
--- a/tests/implicit_diff_test.py
+++ b/tests/implicit_diff_test.py
@@ -15,7 +15,6 @@
 import functools
 
 from absl.testing import absltest
-from absl.testing import parameterized
 
 import jax
 import jax.numpy as jnp
diff --git a/tests/isotonic_test.py b/tests/isotonic_test.py
index 0fc027e7..e477dbaf 100644
--- a/tests/isotonic_test.py
+++ b/tests/isotonic_test.py
@@ -19,7 +19,6 @@
 
 import jax
 import jax.numpy as jnp
-import numpy as onp
 
 from jax.test_util import check_grads
 from jaxopt.isotonic import isotonic_l2_pav
diff --git a/tests/linear_operator_test.py b/tests/linear_operator_test.py
index 599b888e..2e24007f 100644
--- a/tests/linear_operator_test.py
+++ b/tests/linear_operator_test.py
@@ -15,7 +15,6 @@
 
 from absl.testing import absltest
 
-import jax
 import jax.numpy as jnp
 import numpy as onp
 
diff --git a/tests/projected_gradient_test.py b/tests/projected_gradient_test.py
index 8737b8ac..4e16164f 100644
--- a/tests/projected_gradient_test.py
+++ b/tests/projected_gradient_test.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import functools
 
 from absl.testing import absltest
 from absl.testing import parameterized
diff --git a/tests/proximal_gradient_test.py b/tests/proximal_gradient_test.py
index 2c0279da..ef393da3 100644
--- a/tests/proximal_gradient_test.py
+++ b/tests/proximal_gradient_test.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import functools
 from typing import Callable
 
 from absl.testing import absltest
@@ -21,16 +20,12 @@
 import jax
 import jax.numpy as jnp
 
-from jaxopt import implicit_diff
 from jaxopt import objective
-from jaxopt import projection
 from jaxopt import prox
 from jaxopt import ProximalGradient
 from jaxopt._src import test_util
-from jaxopt import tree_util as tu
 
 from sklearn import datasets
-from sklearn import preprocessing
 
 
 def make_stepsize_schedule(max_stepsize, n_steps, power=1.0) -> Callable: