Add an internal definition of ArrayTree and use it instead of chex.ArrayTree.

PiperOrigin-RevId: 831123754

Author: vroulet

docs/conf.py

@@ -125,7 +125,7 @@ def _recursive_add_annotations_import():
     'base.Updates': 'optax.Updates',
     'base.OptState': 'optax.OptState',
     'base.PyTree': 'optax.PyTree',
-    'chex.ArrayTree': 'chex.ArrayTree',
+    'base.ArrayTree': 'optax.ArrayTree',
     'jax.typing.ArrayLike': 'jax.typing.ArrayLike'
 }
 

optax/_src/alias_test.py

@@ -20,7 +20,6 @@
 
 from absl.testing import absltest
 from absl.testing import parameterized
-import chex
 import jax
 from jax import flatten_util
 import jax.numpy as jnp
@@ -433,11 +432,11 @@ def test_gradient_accumulation(self, opt_name, opt_kwargs, dtype):
 
 def _run_opt(
     opt: base.GradientTransformationExtraArgs,
-    fun: Callable[[chex.ArrayTree], jnp.ndarray],
-    init_params: chex.ArrayTree,
+    fun: Callable[[base.ArrayTree], jnp.ndarray],
+    init_params: base.ArrayTree,
     maxiter: int = 500,
     tol: float = 1e-3,
-) -> tuple[chex.ArrayTree, base.OptState]:
+) -> tuple[base.ArrayTree, base.OptState]:
   """Run LBFGS solver by iterative calls to grad transform and apply_updates."""
   value_and_grad_fun = jax.value_and_grad(fun)
 

optax/_src/base.py

@@ -15,10 +15,9 @@
 """Base interfaces and datatypes."""
 
 from collections.abc import Callable
-from typing import (Any, NamedTuple, Optional, Protocol, Sequence, Union,
-                    runtime_checkable)
+from typing import (Any, Iterable, Mapping, NamedTuple, Optional, Protocol,
+                    Sequence, Union, runtime_checkable)
 
-import chex
 import jax
 import jax.numpy as jnp
 
@@ -30,9 +29,11 @@
 PyTree = Any
 Shape = Sequence[int]
 PRNGKey = jax.Array
+ArrayTree = Union[
+    jax.typing.ArrayLike, Iterable['ArrayTree'], Mapping[Any, 'ArrayTree']]
 
-OptState = chex.ArrayTree  # States are arbitrary nests of `jnp.ndarrays`.
-Params = chex.ArrayTree  # Parameters are arbitrary nests of `jnp.ndarrays`.
+OptState = ArrayTree  # States are arbitrary nests of `jnp.ndarrays`.
+Params = ArrayTree  # Parameters are arbitrary nests of `jnp.ndarrays`.
 Updates = Params  # Gradient updates are of the same type as parameters.
 
 Schedule = Callable[[jax.typing.ArrayLike], jax.typing.ArrayLike]

optax/_src/factorized.py

@@ -18,7 +18,6 @@
 import dataclasses
 from typing import NamedTuple, Optional
 
-import chex
 import jax
 import jax.numpy as jnp
 import numpy as np
@@ -81,9 +80,9 @@ class FactoredState(NamedTuple):
   """Overall state of the gradient transformation."""
 
   count: jax.typing.ArrayLike  # number of update steps.
-  v_row: chex.ArrayTree  # Tree of factored params.
-  v_col: chex.ArrayTree  # Tree of factored params.
-  v: chex.ArrayTree  # Tree for params where factoring is skipped.
+  v_row: base.ArrayTree  # Tree of factored params.
+  v_col: base.ArrayTree  # Tree of factored params.
+  v: base.ArrayTree  # Tree for params where factoring is skipped.
 
 
 def scale_by_factored_rms(

optax/_src/linear_algebra.py

@@ -19,7 +19,6 @@
 from typing import Optional, Union
 import warnings
 
-import chex
 import jax
 from jax import lax
 import jax.numpy as jnp
@@ -61,14 +60,14 @@ def _power_iteration_cond_fun(error_tolerance, num_iters, loop_vars):
 
 def power_iteration(
     matrix: Union[
-        jax.typing.ArrayLike, Callable[[chex.ArrayTree], chex.ArrayTree]],
+        jax.typing.ArrayLike, Callable[[base.ArrayTree], base.ArrayTree]],
     *,
-    v0: Optional[chex.ArrayTree] = None,
+    v0: Optional[base.ArrayTree] = None,
     num_iters: jax.typing.ArrayLike = 100,
     error_tolerance: jax.typing.ArrayLike = 1e-6,
     precision: lax.Precision = lax.Precision.HIGHEST,
     key: Optional[base.PRNGKey] = None,
-) -> tuple[jax.typing.ArrayLike, chex.ArrayTree]:
+) -> tuple[jax.typing.ArrayLike, base.ArrayTree]:
   r"""Power iteration algorithm.
 
   This algorithm computes the dominant eigenvalue (i.e. the spectral radius) and

optax/_src/transform.py

@@ -17,7 +17,6 @@
 import functools
 from typing import NamedTuple, Optional
 
-import chex
 import jax
 from jax import nn
 import jax.numpy as jnp
@@ -1504,15 +1503,15 @@ class ScaleByLBFGSState(NamedTuple):
   count: jax.typing.ArrayLike
   params: base.Params
   updates: base.Params
-  diff_params_memory: chex.ArrayTree
-  diff_updates_memory: chex.ArrayTree
+  diff_params_memory: base.ArrayTree
+  diff_updates_memory: base.ArrayTree
   weights_memory: jax.typing.ArrayLike
 
 
 def _precondition_by_lbfgs(
     updates: base.Updates,
-    diff_params_memory: chex.ArrayTree,
-    diff_updates_memory: chex.ArrayTree,
+    diff_params_memory: base.ArrayTree,
+    diff_updates_memory: base.ArrayTree,
     weights_memory: jax.typing.ArrayLike,
     identity_scale: jax.typing.ArrayLike,  # float
     memory_idx: jax.typing.ArrayLike,  # int
@@ -1822,8 +1821,8 @@ def update_fn(
     warn_deprecated_function, replacement='optax.tree.cast'
 )
 def cast_tree(
-    tree: chex.ArrayTree, dtype: Optional[jax.typing.DTypeLike]
-) -> chex.ArrayTree:
+    tree: base.ArrayTree, dtype: Optional[jax.typing.DTypeLike]
+) -> base.ArrayTree:
   return optax.tree.cast(tree, dtype)
 
 

optax/_src/utils.py

@@ -18,7 +18,6 @@
 import functools
 from typing import Optional, Sequence
 
-import chex
 import jax
 import jax.numpy as jnp
 import jax.scipy.stats.norm as multivariate_normal
@@ -57,8 +56,8 @@ def canonicalize_key(key_or_seed: jax.Array | int) -> jax.Array:
     warn_deprecated_function, replacement='optax.tree.cast'
 )
 def cast_tree(
-    tree: chex.ArrayTree, dtype: Optional[jax.typing.DTypeLike]
-) -> chex.ArrayTree:
+    tree: base.ArrayTree, dtype: Optional[jax.typing.DTypeLike]
+) -> base.ArrayTree:
   return optax.tree.cast(tree, dtype)
 
 
@@ -171,29 +170,29 @@ def multi_normal(
 
 @jax.custom_vjp
 def _scale_gradient(
-    inputs: chex.ArrayTree, scale: jax.typing.ArrayLike) -> chex.ArrayTree:
+    inputs: base.ArrayTree, scale: jax.typing.ArrayLike) -> base.ArrayTree:
   """Internal gradient scaling implementation."""
   del scale  # Only used for the backward pass defined in _scale_gradient_bwd.
   return inputs
 
 
 def _scale_gradient_fwd(
-    inputs: chex.ArrayTree, scale: jax.typing.ArrayLike
-) -> tuple[chex.ArrayTree, jax.typing.ArrayLike]:
+    inputs: base.ArrayTree, scale: jax.typing.ArrayLike
+) -> tuple[base.ArrayTree, jax.typing.ArrayLike]:
   return _scale_gradient(inputs, scale), scale
 
 
 def _scale_gradient_bwd(
-    scale: jax.typing.ArrayLike, g: chex.ArrayTree
-) -> tuple[chex.ArrayTree, None]:
+    scale: jax.typing.ArrayLike, g: base.ArrayTree
+) -> tuple[base.ArrayTree, None]:
   return (jax.tree.map(lambda g_: g_ * scale, g), None)
 
 
 _scale_gradient.defvjp(_scale_gradient_fwd, _scale_gradient_bwd)
 
 
 def scale_gradient(
-    inputs: chex.ArrayTree, scale: jax.typing.ArrayLike) -> chex.ArrayTree:
+    inputs: base.ArrayTree, scale: jax.typing.ArrayLike) -> base.ArrayTree:
   """Scales gradients for the backwards pass.
 
   Args:

optax/contrib/_dog.py

@@ -25,7 +25,6 @@
 from collections.abc import Callable
 from typing import Any, NamedTuple, Optional, Union, Literal
 
-import chex
 import jax
 import jax.numpy as jnp
 from optax._src import base
@@ -38,7 +37,7 @@ class DoGState(NamedTuple):
   """State for DoG optimizer."""
 
   is_init_step: jax.Array  # bool
-  init_params: chex.ArrayTree
+  init_params: base.ArrayTree
   max_dist: jax.Array
   sum_sq_norm_grads: jax.Array
 
@@ -220,7 +219,7 @@ def dog(
 class DoWGState(NamedTuple):
   """State for DoWG optimizer."""
 
-  init_params: chex.ArrayTree
+  init_params: base.ArrayTree
   weighted_sq_norm_grads: jax.Array
   estim_sq_dist: jax.Array
 

optax/experimental/_aggregating.py

@@ -17,7 +17,6 @@
 import math
 from typing import Any, NamedTuple, Protocol, Sequence
 
-import chex
 import jax
 import jax.numpy as jnp
 from optax import tree
@@ -26,8 +25,8 @@
 from optax.transforms import _combining
 
 
-PerElementUpdates = chex.ArrayTree
-AggregatedUpdates = chex.ArrayTree
+PerElementUpdates = base.ArrayTree
+AggregatedUpdates = base.ArrayTree
 MaybeAxis = int | Sequence[int] | None
 
 

optax/experimental/_microbatching.py

@@ -20,9 +20,9 @@
 import functools
 from typing import Any, Callable, Sequence, TypeAlias
 
-import chex
 import jax
 import jax.numpy as jnp
+from optax._src import base
 
 
 AccumulatorTree: TypeAlias = Any
@@ -53,10 +53,10 @@ class Accumulator:
       per-microbatch values into a single value. Used by `gvmap`.
   """
 
-  init: Callable[[chex.ArrayTree], chex.ArrayTree]
-  update: Callable[[chex.ArrayTree, chex.ArrayTree, int], chex.ArrayTree]
-  finalize: Callable[[chex.ArrayTree], chex.ArrayTree]
-  aggregate: Callable[[chex.ArrayTree], chex.ArrayTree]
+  init: Callable[[base.ArrayTree], base.ArrayTree]
+  update: Callable[[base.ArrayTree, base.ArrayTree, int], base.ArrayTree]
+  finalize: Callable[[base.ArrayTree], base.ArrayTree]
+  aggregate: Callable[[base.ArrayTree], base.ArrayTree]
 
 
 def _with_floating_check(fn: Callable[..., Any]) -> Callable[..., Any]: