
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

**Summary**

.. autosummary::
  :toctree: _autosummary

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

