reference/generated/tensordict.TensorDictBase.scatter

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.scatter

TensorDictBase.scatter(src: int, tensordicts: list | None = None, *, group: 'torch.distributed.ProcessGroup' | None = None, device: device | str | None = None) → Self

Scatters a list of tensordicts from src to all ranks.

On the source rank, tensordicts must be a list of tensordicts with one element per rank. Each is consolidated, and its metadata + storage are scattered. Non-source ranks receive and reconstruct their tensordict.

Parameters:

• src (int) – The rank of the source process.

• tensordicts (list[TensorDict] | None) – On the source rank, a list of tensordicts to scatter (one per rank). Ignored on other ranks.

Keyword Arguments:

• group (torch.distributed.ProcessGroup, optional) – The process group to use. Defaults to None (default group).

• device (device or str, optional) – The device on which to allocate the receive buffer. If None, the device is inferred from the source’s consolidated storage (transmitted via metadata). Defaults to None.

Returns:

The tensordict assigned to this rank.

Return type:

TensorDict

Every tensordict class and every tensorclass has this method. For the other methods, see Methods by task.