SPMD
----------------------

.. automodule:: flax.linen.spmd
.. currentmodule:: flax.linen

.. autofunction:: Partitioned
.. autofunction:: with_partitioning
.. autofunction:: get_partition_spec
.. autofunction:: get_sharding
.. autofunction:: LogicallyPartitioned
.. autofunction:: logical_axis_rules
.. autofunction:: set_logical_axis_rules
.. autofunction:: get_logical_axis_rules
.. autofunction:: logical_to_mesh_axes
.. autofunction:: logical_to_mesh
.. autofunction:: logical_to_mesh_sharding
.. autofunction:: with_logical_constraint
.. autofunction:: with_logical_partitioning