Flax.linen vs netket.nn for complex models—conflicting information between documentation

[wttai004]

I have seen conflicting advice on whether to use import flax.linen as nn or import netket.nn as nn if my network involves complex numbers.

On the documentation page for Variational State, there is this warning which suggests using the repackaged netket.nn if my model involves complex parameters:

Flax modules do not always work well with complex numbers. For that reason, if your model has complex parameters we invite you to use netket.nn in place of flax.linen and flax.nn. netket.nn re-exports some Flax functionality ensuring that it works nicely with complex numbers. If you don’t find something from Flax in netket.nn, please open an issue in GitHub.

However, when I try to replace the import of nn on my local installation (version 3.13.0), I get this warning message when I initialize some codes. It seems in the source code, the nk.nn.Module and nk.nn.compact is deprecated:

/home1/wttai/.conda/envs/netket/lib/python3.12/site-packages/netket/nn/init.py:90: FutureWarning:

      =======================================================================
      `nk.nn.Module` is deprecated. Use `flax.linen.Module` directly instead.
      =======================================================================

If you imported flax.linen as nn, as is customary to do, you can use nn.Module directly. There are no functionality changes.

_warn_deprecation(
/home1/wttai/.conda/envs/netket/lib/python3.12/site-packages/netket/nn/init.py:90: FutureWarning:

      =======================================================================
      `nk.nn.compact` is deprecated. Use `flax.linen.compact` directly instead.
      =======================================================================

If you imported flax.linen as nn, as is customary to do, you can use nn.compact directly. There are no functionality changes.

_warn_deprecation(

What is the current best practice for complex models? Should I use flax.linen or netket.nn?

[PhilipVinc]

Flax modules do not always work well with complex numbers. For that reason, if your model has complex parameters we invite you to use netket.nn in place of flax.linen and flax.nn. netket.nn re-exports some Flax functionality ensuring that it works nicely with complex numbers. If you don’t find something from Flax in netket.nn, please open an issue in GitHub.

This comment has been outdated since a few years. We upstreamed all of our changes to flax, so you can safely use flax itself.

[PhilipVinc]

(we should remove that warning from the documentation)