update jax.tree.map to comply with jax 0.4.34

[a1302z]

Pull Request to fix #1093

[google-cla]

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.

[vroulet]

Thank you for doing that!
Would it be possible to add a generic test for such functions?
(tree_update_moment, tree_update_infinity_moment and tree_update_moment_per_elem_norm follow the same fix).
Thanks again!

[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.

[vroulet]

Thank you very much, I left you some additional comments

[vroulet]

That's perfect, thank you again @a1302z !