Using optax with partially complex models

Hello!

I’m trying to implement the Fourier Neural Operator model from https://github.com/devzhk/PINO/, translating the Pytorch code there into jax. I ran into an issue where the optax optimizer_state becomes complex for all weight matrices of my model, rather than only the leaves corresponding to complex weights in the model. This only happens after a while, and I’ve only been able to reproduce this while using an optax scheduler. Using only optax.adam seems to work without issue.

Full reproducing script: https://colab.research.google.com/drive/1a5BFz0G20t_VvXOfQqJfqHo4-HL3Ok3X?usp=sharing.

EDIT: This is related to (but possibly different than) https://github.com/deepmind/optax/issues/196

[PhilipVinc]

optax should at least output some form of warning when handling complex parameters…

Yeah, That’s what we tried to argue with optax developers. Hopefully once #297 is merged this will be addressed.

I think my initial observation was due to the gradient clipping. The code for computing the global norm for clipping has the same issue as Adam, and squares a potentially complex number, rather than taking the squared module.

@wdphy16 does PR #279 address gradient clipping as well?

[PhilipVinc]

By the way, can you try using master? The current release of optax does not include the changes to work with complex numbers.

Current master silently computes the wrong norm in adam.