oikosohn's blog

Problem AttributeError: module 'jax' has no attribute 'tree_multimap' Solution Use  jax.tree_util.tree_map  or  jax.tree_map  instead of jax.tree_multimap. JAX 0.3.16  - jax.tree_uti…