jax.sharding.rst

jax.sharding module

.. automodule:: jax.sharding

Classes

.. currentmodule:: jax.sharding

.. autoclass:: Sharding
   :members:

.. autoclass:: SingleDeviceSharding
   :members:
   :show-inheritance:

.. autoclass:: NamedSharding
   :members:
   :show-inheritance:

.. autoclass:: PositionalSharding
   :members:
   :show-inheritance:

.. autoclass:: PmapSharding
   :members:
   :show-inheritance:

.. autoclass:: GSPMDSharding
   :members:
   :show-inheritance:

.. autoclass:: PartitionSpec
   :members:

.. autoclass:: Mesh
   :members: