Update jax.tree_util.tree_map to jax.tree.map#1821

Merged

fehiepsi merged 5 commits intopyro-ppl:masterfrom fehiepsi:jaxtree

Jun 26, 2024

+236-243

Commits on Jun 25, 2024

also fix deprecation warning of using a_min, a_max in jnp.clip
fehiepsi
committed