Better tests for utils

[acforvs]

Hi, hopefully this one closes #404

[hbq1]

Hi @acforvs, many thanks for your PR and apologies for the belated review!
It looks good to me; I left only a few comments. Could you also add tests for the rest of helpers from utils.py, so we could close the original issue?

[acforvs]

Hi @hbq1 , thank you for the review!

Sure thing, I'll add the tests for the remaining functions, but I have some questions about them:

1. The docstring for set_diags states that the shape of new_diags should be NxD. However, new_diags is only used in one place where it is flattened.

As a result, examples like

a = jnp.ones(shape=(3, 2, 2))
diags = jnp.array([[[3, 3, 3, 3, 3, 3]]])
print(a.shape, diags.shape)  # prints (3, 2, 2) (1, 1, 6)
optax._src.utils.set_diags(a, diags)

or

a = jnp.ones(shape=(3, 2, 2))
diags = jnp.array([[3, 3, 3], [3, 3, 3]])
print(a.shape, diags.shape)  # prints (3, 2, 2) (2, 3)
optax._src.utils.set_diags(a, diags)

work just fine

Is this by design or should we raise an error in case of a shape mismatch?
Currently, the function is used internally in two places: 1, 2. The comments specify that the dimension of the passed variable is indeed NxD there, so I'd suggest adding another assert. What are your thoughts?

2. cast_tree uses .astype which wouldn't work for simple datatypes, such as

tree = dict(a=2.5, b=dict(c=-2.5))
tree = jax.tree_util.tree_map(lambda x : x, tree)
optax._src.utils.cast_tree(tree, int)

Would it be okay to only test this function for dicts of jnp.arrays of differerent dtypes? Something like

tree = dict(a=jnp.array(2.5), b=dict(c=jnp.array(-2.5)))
tree = jax.tree_util.tree_map(lambda x : x, tree)
optax._src.utils.cast_tree(tree, int)

[hbq1]

Great, thank you!

1. It is used in some internal code, so the correct behaviour would be to raise an exception in case of shape mismatch; it'd be great if you add this to the implementation!
2. Testing only with trees of jnp.arrays would suffice.

[acforvs]

I've fixed the errors and added tests for the rest of the utils.py functions.
I also changed assert to a ValueError in set_diags. Please let me know how it looks

[hbq1]

Thanks for the changes and for the nice PR 👍