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
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
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:
![swh spinner](/static/img/swh-spinner.gif)
Computing file changes ...