Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Deprecate jax.tree_map in favor of jax.tree.map #19930

Merged
merged 1 commit into from
Feb 22, 2024

Conversation

jakevdp
Copy link
Collaborator

@jakevdp jakevdp commented Feb 22, 2024

No description provided.

@jakevdp jakevdp self-assigned this Feb 22, 2024
@yashk2810
Copy link
Member

Why are we deprecating it? I thought the plan was too support both?

@hawkinsp
Copy link
Member

Why are we deprecating it? I thought the plan was too support both?

Nope, the plan is to just have jax.tree.*.

@jakevdp
Copy link
Collaborator Author

jakevdp commented Feb 22, 2024

Why are we deprecating it? I thought the plan was too support both?

The plan is to support both jax.tree.map and jax.tree_util.tree_map, but remove tree-related methods from the top-level jax namespace. Other jax.tree_* methods have already been deprecated; jax.tree_map was a holdout because it's very common and users didn't like the more verbose jax.tree_util.tree_map. Now that jax.tree.map is an option, we can deprecate jax.tree_map.

@jakevdp jakevdp added the pull ready Ready for copybara import and testing label Feb 22, 2024
@copybara-service copybara-service bot merged commit be002b5 into google:main Feb 22, 2024
13 checks passed
@jakevdp jakevdp deleted the dep-tree_map branch February 23, 2024 00:13
copybara-service bot pushed a commit to google/CommonLoopUtils that referenced this pull request Apr 4, 2024
…https://jax.readthedocs.io/en/latest/changelog.html#jax-0-4-26-april-3-2024), `jax.tree_map` is deprecated. We should switch to `jax.tree.map` if/when the JAX minimum version is >=0.4.26.

Additional context:
* JAX PR that added the `jax.tree` module: google/jax#19588
* JAX PR that deprecated `jax.tree_map`: google/jax#19930

PiperOrigin-RevId: 621905494
copybara-service bot pushed a commit to google/CommonLoopUtils that referenced this pull request Apr 4, 2024
…https://jax.readthedocs.io/en/latest/changelog.html#jax-0-4-26-april-3-2024), `jax.tree_map` is deprecated. We should switch to `jax.tree.map` if/when the JAX minimum version is >=0.4.26 for CLU.

Additional context:
* JAX PR that added the `jax.tree` module: google/jax#19588
* JAX PR that deprecated `jax.tree_map`: google/jax#19930

PiperOrigin-RevId: 621905494
copybara-service bot pushed a commit to google/CommonLoopUtils that referenced this pull request Apr 5, 2024
…https://jax.readthedocs.io/en/latest/changelog.html#jax-0-4-26-april-3-2024), `jax.tree_map` is deprecated. We should switch to `jax.tree.map` if/when the JAX minimum version is >=0.4.26 for CLU.

Additional context:
* JAX PR that added the `jax.tree` module: google/jax#19588
* JAX PR that deprecated `jax.tree_map`: google/jax#19930

PiperOrigin-RevId: 621905494
copybara-service bot pushed a commit to google/CommonLoopUtils that referenced this pull request Apr 5, 2024
…https://jax.readthedocs.io/en/latest/changelog.html#jax-0-4-26-april-3-2024), `jax.tree_map` is deprecated. We should switch to `jax.tree.map` if/when the JAX minimum version is >=0.4.26 for CLU.

Additional context:
* JAX PR that added the `jax.tree` module: google/jax#19588
* JAX PR that deprecated `jax.tree_map`: google/jax#19930

PiperOrigin-RevId: 622092167
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
pull ready Ready for copybara import and testing
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

3 participants