jax.sharding module

.. automodule:: jax.sharding

Classes

.. currentmodule:: jax.sharding

.. autoclass:: Sharding :members: .. autoclass:: SingleDeviceSharding :members: :show-inheritance: .. autoclass:: NamedSharding :members: :show-inheritance: .. autoclass:: PartitionSpec :members: .. autoclass:: Mesh :members: