Fix dtype casting inside tree_add_scale.

[carlosgmartin]

The current behavior of tree_add_scale is confusing:

$ py -c "import optax; print(optax.tree_utils.tree_add_scale([1], 0.1, [1]))"
[Array(1, dtype=int32)]

This is because its current implementation casts scalar to the dtype of x.

This PR modifies tree_add_scale to remove this casting and thus get the expected behavior:

$ py -c "import optax; print(optax.tree_utils.tree_add_scale([1], 0.1, [1]))"
[Array(1.1, dtype=float32, weak_type=True)]

Subsequently, it fixes the scan loop bodies inside _precondition_by_lbfgs to use tree_cast_like to ensure the carry output matches the carry input.

[vroulet]

@mtthss you introduced this dtype conversion in #926. Do you remember why?

@carlosgmartin add a test. It's clearly a bug so a test is welcome. Thanks again!

[vroulet]

Also @carlosgmartin, can you explain

1. the bug that you then get with lbfgs
2. whether it was caught by a test or whether you found it yourself

There may be other functions that work well thanks to this undocumented dtype conversion. So maybe we'll have to make a pass on references to tree_add_scale

[carlosgmartin]

@vroulet The induced LBFGS bug was caught by a test. Without the two

vec_new = optax.tree.cast_like(vec_new, vec)

lines,

$ pytest optax/_src/alias_test.py -k test_state_shape_dtype_shard_stability

outputs

optax/_src/alias_test.py ..s..........s..ss...........ss...........ss.....ss [ 25%]
.................s......ss..s........ss.....ss.............................. [ 62%]
..........................ss...........ss..................ss....FF..FF..... [ 99%]
.                                                                            [100%]

===================================== FAILURES =====================================
________________ LBFGSTest.test_state_shape_dtype_shard_stability0 _________________

self = <alias_test.LBFGSTest testMethod=test_state_shape_dtype_shard_stability0>
dtype = <class 'jax.numpy.float16'>
linesearch = GradientTransformationExtraArgs(init=<function scale_by_backtracking_linesearch.<locals>.init_fn at 0x127a46980>, update=<function scale_by_backtracking_linesearch.<locals>.update_fn at 0x127a46b60>)

...

E       TypeError: scan body function carry input and carry output must have equal types, but they differ:
E       
E       The input carry vec has type float16[3] but the corresponding output carry component has type float32[3], so the dtypes do not match.
E       
E       Revise the function so that all output types match the corresponding input types.

...

[rdyro]

Remove this test please

[carlosgmartin]

@vroulet asked me to add a test here. Can you clarify?

[rdyro]

Removing the dtype cast in tree_add_scale is a good idea and making this explicit for L-BFGS is great! Thanks @carlosgmartin