update jax.tree.map to comply with jax 0.4.34

[a1302z]

Pull Request to fix #1093

[google-cla[bot]]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[a1302z]

Happy to do that. How would you imagine this test to be fully generic? A simple test case would be something like this (extended to all corresponding functions):

  def test_none_arguments(self):
    m = jnp.logspace(-10, 10, num=21, dtype=jnp.bfloat16)  # 1e-10 ... 1e10
    tu.tree_update_moment(None, m, 0.9, 1)

Let me know what you think.

[vroulet]

Yes, I would even simplify it by using just jnp.array([1., 2., 3.]) instead of jnp.logspace. Thank you again for doing that!

[a1302z]

I have added test cases for tree_add_scalar_mul, tree_update_moment, tree_update_infinity_moment, tree_update_moment_per_elem_norm, apply_updates, incremental_update, keep_params_nonnegative, and add_decayed_weights. I have omitted scale_by_adam as mu would need to be None, which is not an input argument to the function.