search

brentyi / jax_dataclasses

Pytrees + dataclasses ❤️
MIT License
61 stars 6 forks source link

`jax.tree_leaves` is deprecated #4

Closed lucagrementieri closed 2 years ago

lucagrementieri commented 2 years ago

The file jax_dataclasses/_copy_and_mutate.py raises many warnings complaining a deprecated function.

FutureWarning: jax.tree_leaves is deprecated, and will be removed in a future release. Use jax.tree_util.tree_leaves instead.
brentyi commented 2 years ago

Thanks! Should be fixed in the latest release; let me know if you're still having issues. :slightly_smiling_face: