Skip to main content

Browse the archive

Revision 01194bd2fb1f7b6a46db88be14324f7acdb1723e authored by Sergei Lebedev on 16 May 2024, 18:27:56 UTC, committed by jax authors on 16 May 2024, 18:29:09 UTC 
Clarified the type of the inputs to callback APIs
 
The callback APIs were migrated to use jax.Arrays for both inputs and outputs
in JAX 0.4.27.

PiperOrigin-RevId: 634473890
1 parent 380503b
Raw File
jax.sharding.rst 
``jax.sharding`` module
=======================

.. automodule:: jax.sharding

Classes
-------

.. currentmodule:: jax.sharding

.. autoclass:: Sharding
   :members:
.. autoclass:: XLACompatibleSharding
   :members:
   :show-inheritance:
.. 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:
back to top