Implement split real norm

[wdphy16]

As proposed in #196, it is an optimizer wrapper that splits the complex parameters into pairs of real parameters before sending them to the update of the wrapped optimizer, and merges the pairs of real updates into complex updates afterward.

[mkunesch]

Hi! Thanks a lot for the PR!
(Assigning myself - I have started reviewing and should have time to finish tomorrow.)

[mkunesch]

Hey!

Thank you so much for this PR and your detailed explanation in the gist attached to #196. This was super helpful for the review.

I think it looks great! I just have a few very minor comments and nits.

Thanks a lot again and let me know if you have any questions regarding the comments!

[wdphy16]

Hi @mkunesch , thanks for your detailed review! I've made the changes accordingly.

[mkunesch]

Looks great! Thanks for making the changes!

[rosshemsley]

Thanks! It's really great to see these additions :)

I left a bunch of thoughts to compliment @mkunesch 's review - I'll let him take over in managing getting this thing merged, though :)

[mkunesch]

Thanks a lot @rosshemsley for the additional comments and @wdphy16 for implementing them!

[wdphy16]

I've added the experimental namespace by mimicking haiku.experimental. Please see if you like it.

pylint reports that there are unused imports in experimental.py. How do you usually skip that?

[mkunesch]

Hi! Thanks a lot! It looks great to me! (Only flagged a missing full stop above)

Re the import error: I mean, we can easily disable it (for __init__.py it's disabled in the pylint settings afaik) but I'm also not sure why haiku adds a file for experimental.py rather than adding an experimental folder in _src and importing it in __init__.py. I assume that would work too and we wouldn't have to disable warnings then? (Also tagging @rosshemsley for an opinion on the structure).

[rosshemsley]

+1 to @mkunesch - I would go for

haiku/
   _src/
      experimental/
          complex.py

    __init__.py
        experimental/
              complex.py (or __init__.py)

[wdphy16]

Ok, this project structure is also more intuitive to me

[mkunesch]

Thanks a lot! Looks great to me. I like the folder structure and it passes the tests/lint checks.
Thanks a lot for the contribution again and for making the changes!

[wdphy16]

Hi @mkunesch , shall we merge this now?

[mkunesch]

Hi! This got merged yesterday :-D! (Sorry it took slightly longer - it was blocked by a flaky test d663336).

Thank you very much for this PR again and for all your work (especially the excellent design doc!) on how to support complex numbers in optax!