mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 03:46:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Re-land #11498 after internal fixes.

Browse Source 
maintain an alias to `jax.tree_util.tree_map` in the top level `jax` module

PiperOrigin-RevId: 463885774
This commit is contained in:
Lena Martens 2022-07-28 11:33:01 -07:00 committed by jax authors
parent 697c9b1736
commit 8ca5ecc7f3
2 changed files with 5 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
CHANGELOG.md
View File

@ -31,12 +31,12 @@ PLEASE REMEMBER TO CHANGE THE '..main' WITH AN ACTUAL TAG in GITHUB LINK.
* Binary operations between JAX arrays and built-in collections (`dict`, `list`, `set`, `tuple`)
now raise a `TypeError` in all cases. Previously some cases (particularly equality and inequality)
would return boolean scalars inconsistent with similar operations in NumPy ({jax-issue}`#11234`).
* {mod}`jax.tree_util` routines accessed as top-level JAX package imports are now deprecated, and
will be removed in a future JAX release in accordance with the {ref}`api-compatibility` policy:
* Several {mod}`jax.tree_util` routines accessed as top-level JAX package imports are now
deprecated, and will be removed in a future JAX release in accordance with the
{ref}`api-compatibility` policy:
* {func}`jax.treedef_is_leaf` is deprecated in favor of {func}`jax.tree_util.treedef_is_leaf`
* {func}`jax.tree_flatten` is deprecated in favor of {func}`jax.tree_util.tree_flatten`
* {func}`jax.tree_leaves` is deprecated in favor of {func}`jax.tree_util.tree_leaves`
* {func}`jax.tree_map` is deprecated in favor of {func}`jax.tree_util.tree_map`
* {func}`jax.tree_structure` is deprecated in favor of {func}`jax.tree_util.tree_structure`
* {func}`jax.tree_transpose` is deprecated in favor of {func}`jax.tree_util.tree_transpose`
* {func}`jax.tree_unflatten` is deprecated in favor of {func}`jax.tree_util.tree_unflatten`

4
jax/__init__.py
View File

@ -114,12 +114,12 @@ from jax.experimental.maps import soft_pmap as soft_pmap
from jax.version import __version__ as __version__
from jax.version import __version_info__ as __version_info__
# TODO(jakevdp): remove these deprecated routines after October 2022
from jax._src.tree_util import (
tree_map as tree_map,
# TODO(jakevdp): remove these deprecated routines after October 2022
_deprecated_treedef_is_leaf as treedef_is_leaf,
_deprecated_tree_flatten as tree_flatten,
_deprecated_tree_leaves as tree_leaves,
_deprecated_tree_map as tree_map,
_deprecated_tree_structure as tree_structure,
_deprecated_tree_transpose as tree_transpose,
_deprecated_tree_unflatten as tree_unflatten,