Alex Ziller

This works wonderfully. Instead of cloning the entire dataset now, I just have links to the original image, which saves so much time and memory. I am very much in...

I found it by reproducing this [tutorial of equinox](https://docs.kidger.site/equinox/examples/mnist/). I found eight more occurrences and have modified them appropriately. I'll create a PR :)

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

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