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

Copybara import of the project:

Browse Source 
--
e1f1e93e0c8b53e62a064b06b56c84a2bfedb911 by Roy Frostig <frostig@google.com>:

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

PiperOrigin-RevId: 461146464
This commit is contained in:
jax authors 2022-07-15 01:23:27 -07:00
parent e19f6973c4
commit 023e6f5955
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

@ -17,12 +17,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`).
* 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:
* {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

@ -115,12 +115,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_multimap as tree_multimap,
_deprecated_tree_structure as tree_structure,
_deprecated_tree_transpose as tree_transpose,