Markus Kunesch

Hi! Interesting - thanks for reporting this! Are you also at more than ~2/3 memory usage when you use `apply_every`? From a first look, I could see that the implementation...

Hi! I think the example looks great in terms of demonstrating how to use `multi_transform` with `set_to_zero`. In this particular example, since sigma is never optimized it might be worth...

Interesting, based on your description this would only happen if the dtype inference in line [383](https://github.com/deepmind/optax/blob/233294c60034bb3852c966601a59366ee45adfd9/optax/_src/wrappers.py#L383) results in the wrong type so I could try looking into whether the dtype...

Interesting, I think that would be great! Thanks a lot! Let us know if you'd like to discuss anything about the implementation as you write it. How much of the...

Hi, thanks a lot for pointing this out! This is definitely something we should discuss especially if it would be convenient for flax to have these types supported. I think...

That's great, thanks a lot! The replies in #206 might be helpful for showing the different ways of doing it. Let us know if you have any questions!