Remove the deprecation of jax.tree_map for the release of 0.4.25

PiperOrigin-RevId: 610014256

CHANGELOG.md

@@ -15,8 +15,6 @@ Remember to align the itemized text with the first line of an item within a list
  * JAX arrays now support NumPy-style scalar boolean indexing, e.g. `x[True]` or `x[False]`.
  * Added {mod}`jax.tree` module, with a more convenient interface for referencing functions
  in {mod}`jax.tree_util`.
- * {func}`jax.tree_map` is deprecated; use `jax.tree.map` instead, or for backward
- compatibility with older JAX versions, use {func}`jax.tree_util.tree_map`.
  * {func}`jax.tree.transpose` (i.e. {func}`jax.tree_util.tree_transpose`) now accepts
  `inner_treedef=None`, in which case the inner treedef will be automatically inferred.
 

jax/__init__.py

@@ -136,7 +136,7 @@
 )
 
 from jax._src.tree_util import (
- tree_map as _deprecated_tree_map,
+ tree_map as tree_map,
  treedef_is_leaf as _deprecated_treedef_is_leaf,
  tree_flatten as _deprecated_tree_flatten,
  tree_leaves as _deprecated_tree_leaves,
@@ -212,20 +212,13 @@
  "or jax.tree_util.tree_unflatten (any JAX version).",
  _deprecated_tree_unflatten
  ),
- # Added Feb 22, 2024
- "tree_map": (
- "jax.tree_map is deprecated: use jax.tree.map (jax v0.4.25 or newer) "
- "or jax.tree_util.tree_map (any JAX version).",
- _deprecated_tree_map
- ),
 }
 
 import typing as _typing
 if _typing.TYPE_CHECKING:
  from jax._src.tree_util import treedef_is_leaf as treedef_is_leaf
  from jax._src.tree_util import tree_flatten as tree_flatten
  from jax._src.tree_util import tree_leaves as tree_leaves
- from jax._src.tree_util import tree_map as tree_map
  from jax._src.tree_util import tree_structure as tree_structure
  from jax._src.tree_util import tree_transpose as tree_transpose
  from jax._src.tree_util import tree_unflatten as tree_unflatten