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.

Every rank of the process group calls it, in an initialized torch.distributed process group. On the source rank, tensordicts holds one tensordict per rank; each is consolidated (see consolidate()) and sent as one buffer. The other ranks rebuild their tensordict from the buffer they receive. The tensordict the method is called on is not used. See Distributed tensordicts.

Parameters:

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

• tensordicts (list of TensorDictBase, optional) – on the source rank, the tensordicts to scatter, one per rank. Ignored on the other ranks. Defaults to None.

Keyword Arguments:

• group (torch.distributed.ProcessGroup, optional) – the process group. Defaults to None, which uses the default group.

• device (device or str, optional) – on the receiving ranks, the device of the receive buffer. Defaults to None, which uses the device of the consolidated storage of the source.

Returns:

The tensordict that this rank receives.

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