rocm_jax/docs/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:
DOC: make API doc titles more uniform 2023-01-18 10:59:42 -08:00			``jax.sharding`` module
			`=======================`
Add docstrings for `Sharding` classes. Right now I am only documenting `Sharding`, `XLACompatibleSharding`, `MeshPspecSharding` and `SingleDeviceSharding`. Also moving jax_array_migration guide to reference documentation. PiperOrigin-RevId: 488489503 2022-11-14 15:47:06 -08:00
			`.. automodule:: jax.sharding`

			`Classes`
			`-------`

			`.. currentmodule:: jax.sharding`

			`.. autoclass:: Sharding`
			`:members:`
Improve the sharding documentation. * do some proofreading. * add PmapSharding and GSPMDSharding, which are both missing. 2023-08-03 10:15:09 -04:00			`.. autoclass:: SingleDeviceSharding`
			`:members:`
			`:show-inheritance:`
Add docstrings for `Sharding` classes. Right now I am only documenting `Sharding`, `XLACompatibleSharding`, `MeshPspecSharding` and `SingleDeviceSharding`. Also moving jax_array_migration guide to reference documentation. PiperOrigin-RevId: 488489503 2022-11-14 15:47:06 -08:00			`.. autoclass:: NamedSharding`
			`:members:`
			`:show-inheritance:`
Improve the sharding documentation. * do some proofreading. * add PmapSharding and GSPMDSharding, which are both missing. 2023-08-03 10:15:09 -04:00			`.. autoclass:: PositionalSharding`
			`:members:`
			`:show-inheritance:`
			`.. autoclass:: PmapSharding`
			`:members:`
			`:show-inheritance:`
			`.. autoclass:: GSPMDSharding`
Add docstrings for `Sharding` classes. Right now I am only documenting `Sharding`, `XLACompatibleSharding`, `MeshPspecSharding` and `SingleDeviceSharding`. Also moving jax_array_migration guide to reference documentation. PiperOrigin-RevId: 488489503 2022-11-14 15:47:06 -08:00			`:members:`
			`:show-inheritance:`
Merge pull request #14248 from jakevdp:dead-code PiperOrigin-RevId: 506405131 2023-02-01 21:25:46 +00:00			`.. autoclass:: PartitionSpec`
			`:members:`
			`.. autoclass:: Mesh`
			`:members:`