Remove use of cast_tree in favor of optax.tree.cast

PiperOrigin-RevId: 830951266

Author: vroulet

optax/_src/transform.py

@@ -24,7 +24,6 @@
 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
@@ -1812,16 +1811,6 @@ def update_fn(
 
 ### Legacy symbols to be removed. ###
 
-
-@functools.partial(
-    warn_deprecated_function, replacement='optax.tree.cast'
-)
-def cast_tree(
-    tree: chex.ArrayTree, dtype: Optional[chex.ArrayDType]
-) -> chex.ArrayTree:
-  return optax.tree.cast(tree, dtype)
-
-
 trace = _accumulation.trace
 TraceState = _accumulation.TraceState
 ema = _accumulation.ema

optax/_src/utils.py

@@ -27,7 +27,6 @@
 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,15 +54,6 @@ def canonicalize_key(key_or_seed: jax.Array | int) -> jax.Array:
   return jax.random.key(key_or_seed)
 
 
-@functools.partial(
-    warn_deprecated_function, replacement='optax.tree.cast'
-)
-def cast_tree(
-    tree: chex.ArrayTree, dtype: Optional[chex.ArrayDType]
-) -> chex.ArrayTree:
-  return optax.tree.cast(tree, dtype)
-
-
 def set_diags(a: jax.Array, new_diags: chex.Array) -> chex.Array:
   """Set the diagonals of every DxD matrix in an input of shape NxDxD.