spmd
------------------------

.. automodule:: flax.nnx
.. currentmodule:: flax.nnx

.. autofunction:: get_partition_spec
.. autofunction:: get_named_sharding
.. autofunction:: with_partitioning