Use jax.tree_util.tree_map in place of deprecated tree_multimap.
Use jax.tree_util.tree_map in place of deprecated tree_multimap.

The latter is a simple alias of the former, so this change is a no-op.

PiperOrigin-RevId: 461177802

deepmind/optax,deepmind/optax