etha.comm.get_chunks

Build chunk descriptors from routes.

Functions

calculate_chunk_shape(→ tuple[int, Ellipsis])
m2m_to_chunks(→ list[etha.comm.ir.Chunk])	Materialize an M2MMap (topology) onto local tensors into chunks.

Module Contents

etha.comm.get_chunks.calculate_chunk_shape(num_slicers: list[int], tensor_shape: tuple[int, Ellipsis] | None) → tuple[int, Ellipsis]

etha.comm.get_chunks.m2m_to_chunks(m2m: etha.comm.ir.M2MMap, rank: int, source_tensor: torch.Tensor | None = None, target_tensor: torch.Tensor | None = None, transfer_dtype: torch.dtype | None = None, source_partial_groups: list[tuple[torch.distributed.ProcessGroup, str]] | None = None) → list[etha.comm.ir.Chunk]

Materialize an M2MMap (topology) onto local tensors into chunks.