google / google/flax

Remove `tree_map` deprecation filter after Flax upgrades minimum Python version to 3.10

Open
#3,842 0 comments 0 reactions 1 assignee Claimed by @chiamp View on GitHub
Priority: P2 - no schedule
Dominant language
Jupyter Notebook
Stars
7.3k
Forks
833
Avg merge
5h 11m
Merged PRs (30d)
5

Description

Context:
* As of [JAX 0.4.26](https://jax.readthedocs.io/en/latest/changelog.html#jax-0-4-26-april-3-2024), `jax.tree_map` is deprecated
* #3823 renames all `jax.tree_map` usages to `jax.tree_util.tree_map` in Flax, however we get an [error](https://github.com/google/flax/actions/runs/8562140920/job/23464811629?pr=3797#step:9:467) in CI because of a CLU dependency
* After [fixing](https://github.com/google/CommonLoopUtils/pull/342) CLU and pushing a new [release](https://github.com/google/CommonLoopUtils/releases/tag/v0.0.12), the [error](https://github.com/google/flax/actions/runs/8637546849/job/23679980897?pr=3836#step:9:459) remains on CI for Python 3.9 tests
* this is because Flax [enforces an earlier version of CLU](https://github.com/google/flax/blob/main/pyproject.toml#L44) (before the `tree_map` fix) on python versions less than 3.10, since the [match-case](https://github.com/google/CommonLoopUtils/blob/main/clu/parameter_overview.py#L174) syntax used by CLU is only available in Python 3.10 or greater
* i.e. because Flax supports a minimum Python version of 3.9, Flax must use an earlier CLU version (where the CLU fix has not landed yet), since the current CLU version with the fix uses a python syntax that isn't available until Python 3.10
* Our current solution is to add a deprecation warning filter in #3828

Once Flax upgrades its minimum Python version to 3.10, we should remove the [deprecation warning filter](https://github.com/google/flax/blob/main/pyproject.toml#L146) and remove [enforcing](https://github.com/google/flax/blob/main/pyproject.toml#L44) an earlier version of CLU.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.