nki.collectives
===============

.. currentmodule:: nki.collectives


.. _nki-collectives:

NKI Collectives
---------------

Collective operations for multi-rank communication.

.. autosummary::
   :toctree: generated
   :nosignatures:

   all_reduce
   all_gather
   reduce_scatter
   all_to_all
   collective_permute
   collective_permute_implicit
   collective_permute_implicit_reduce
   collective_permute_implicit_current_processing_rank_id
   rank_id


Helper Classes
--------------

.. autosummary::
   :toctree: generated
   :template: nki-custom-class-template.rst
   :nosignatures:

   ReplicaGroup