Use internal warn_deprecated_function instead of the Chex version.

PiperOrigin-RevId: 830563605

Author: emilyfertig

optax/_src/deprecations.py

@@ -75,7 +75,7 @@ def warn_deprecated_function(
     The wrapped function.
 
   Example usage:
-  >>> @functools.partial(chex.warn_deprecated_function, replacement='g')
+  >>> @functools.partial(warn_deprecated_function, replacement='g')
   ... def f(a, b):
   ...   return a + b
   """

optax/_src/transform.py

@@ -24,6 +24,7 @@
 from optax._src import base
 from optax._src import numerics
 from optax._src import utils
+from optax._src.deprecations import warn_deprecated_function  # pylint: disable=g-importing-member
 from optax.transforms import _accumulation
 from optax.transforms import _adding
 import optax.tree
@@ -1813,7 +1814,7 @@ def update_fn(
 
 
 @functools.partial(
-    chex.warn_deprecated_function, replacement='optax.tree.cast'
+    warn_deprecated_function, replacement='optax.tree.cast'
 )
 def cast_tree(
     tree: chex.ArrayTree, dtype: Optional[chex.ArrayDType]

optax/_src/utils.py

@@ -27,6 +27,7 @@
 import jax.scipy.stats.norm as multivariate_normal
 from optax._src import base
 from optax._src import numerics
+from optax._src.deprecations import warn_deprecated_function  # pylint: disable=g-importing-member
 import optax.tree
 
 
@@ -55,7 +56,7 @@ def canonicalize_key(key_or_seed: jax.Array | int) -> jax.Array:
 
 
 @functools.partial(
-    chex.warn_deprecated_function, replacement='optax.tree.cast'
+    warn_deprecated_function, replacement='optax.tree.cast'
 )
 def cast_tree(
     tree: chex.ArrayTree, dtype: Optional[chex.ArrayDType]

optax/losses/_classification.py

@@ -23,6 +23,7 @@
 import jax.numpy as jnp
 import jax.scipy.special
 from optax import projections
+from optax._src.deprecations import warn_deprecated_function  # pylint: disable=g-importing-member
 
 
 def canonicalize_axis(axis, ndim):
@@ -83,7 +84,7 @@ class is an independent binary prediction and different classes are not
 
 
 @functools.partial(
-    chex.warn_deprecated_function, replacement='sigmoid_binary_cross_entropy'
+    warn_deprecated_function, replacement='sigmoid_binary_cross_entropy'
 )
 def binary_logistic_loss(logits, labels):
   return sigmoid_binary_cross_entropy(logits, labels)
@@ -147,7 +148,7 @@ def sparsemax_loss(
   return jax.nn.sparse_plus(jnp.where(labels, -logits, logits))
 
 
-@functools.partial(chex.warn_deprecated_function, replacement='sparsemax_loss')
+@functools.partial(warn_deprecated_function, replacement='sparsemax_loss')
 def binary_sparsemax_loss(logits, labels):
   return sparsemax_loss(logits, labels)
 
@@ -406,7 +407,7 @@ def softmax_cross_entropy_with_integer_labels(
 
 
 @functools.partial(
-    chex.warn_deprecated_function,
+    warn_deprecated_function,
     replacement='softmax_cross_entropy_with_integer_labels',
 )
 def multiclass_logistic_loss(logits, labels):

optax/monte_carlo/control_variates.py

@@ -60,6 +60,7 @@
 import jax
 import jax.numpy as jnp
 from optax._src import base
+from optax._src.deprecations import warn_deprecated_function  # pylint: disable=g-importing-member
 
 
 CvState = Any
@@ -69,7 +70,7 @@
 ControlVariate = tuple[ComputeCv, CvExpectedValue, UpdateCvState]
 
 
-@chex.warn_deprecated_function
+@warn_deprecated_function
 def control_delta_method(
     function: Callable[[chex.Array], float],
 ) -> ControlVariate:
@@ -143,7 +144,7 @@ def update_state(
   return delta, expected_value_delta, update_state
 
 
-@chex.warn_deprecated_function
+@warn_deprecated_function
 def moving_avg_baseline(
     function: Callable[[chex.Array], float],
     decay: float = 0.99,
@@ -217,7 +218,7 @@ def _map(cv, params, samples, state):
   return jax.vmap(lambda x: cv(params, x, state))(samples)
 
 
-@chex.warn_deprecated_function
+@warn_deprecated_function
 def control_variates_jacobians(
     function: Callable[[chex.Array], float],
     control_variate_from_function: Callable[
@@ -366,7 +367,7 @@ def param_fn(x):
   return jacobians, control_variate_state
 
 
-@chex.warn_deprecated_function
+@warn_deprecated_function
 def estimate_control_variate_coefficients(
     function: Callable[[chex.Array], float],
     control_variate_from_function: Callable[

optax/monte_carlo/stochastic_gradient_estimators.py

@@ -38,9 +38,10 @@
 import numpy as np
 from optax._src import base
 from optax._src import utils
+from optax._src.deprecations import warn_deprecated_function  # pylint: disable=g-importing-member
 
 
-@chex.warn_deprecated_function
+@warn_deprecated_function
 def score_function_jacobians(
     function: Callable[[chex.Array], float],
     params: base.Params,
@@ -93,7 +94,7 @@ def surrogate(params):
   return jax.jacfwd(surrogate)(params)
 
 
-@chex.warn_deprecated_function
+@warn_deprecated_function
 def pathwise_jacobians(
     function: Callable[[chex.Array], float],
     params: base.Params,
@@ -146,7 +147,7 @@ def surrogate(params):
   return jax.jacfwd(surrogate)(params)
 
 
-@chex.warn_deprecated_function
+@warn_deprecated_function
 def measure_valued_jacobians(
     function: Callable[[chex.Array], float],
     params: base.Params,
@@ -207,7 +208,7 @@ def measure_valued_jacobians(
   ]
 
 
-@chex.warn_deprecated_function
+@warn_deprecated_function
 def measure_valued_estimation_mean(
     function: Callable[[chex.Array], float],
     dist: Any,
@@ -280,7 +281,7 @@ def measure_valued_estimation_mean(
   return grads
 
 
-@chex.warn_deprecated_function
+@warn_deprecated_function
 def measure_valued_estimation_std(
     function: Callable[[chex.Array], float],
     dist: Any,