reference/generated/tensordict.TensorDictBase.recv

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

TensorDictBase.recv

TensorDictBase.recv(src: int | 'TensorDictPipe' | None = None, *, group: 'torch.distributed.ProcessGroup' | None = None, group_src: int | None = None, init_tag: int = 0, pseudo_rand: bool = False, consolidated: bool = False) → int

Receives the leaves that send() sends, and writes them into the tensordict in place.

The tensordict must have the keys, shapes and dtypes of the tensordict that the sender sends. See send() for an example, and Distributed tensordicts.

Parameters:

src (int or TensorDictPipe, optional) – the global rank of the sender, or a TensorDictPipe to receive from with UCXX, which ignores the other arguments. Exclusive with group_src: pass exactly one of them. Defaults to None.

Keyword Arguments:

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

• group_src (int, optional) – the rank of the sender in group, which must be given. Exclusive with src. The messages then go through the backend of group directly, so group can be a standalone ProcessGroup built on a store, never registered by init_process_group() or new_group(). Defaults to None.

• init_tag (int, optional) – the init_tag that the sender passed. Defaults to 0.

• pseudo_rand (bool, optional) – the pseudo_rand that the sender passed. Defaults to False.

• consolidated (bool, optional) – if True, receives one message into the consolidated storage of the tensordict, which updates all its leaves, as sent by send(consolidated=True). The tensordict must be consolidated already, for example by consolidate(metadata=True) or from_remote_init(); otherwise an AttributeError is raised. init_tag and pseudo_rand are ignored. Defaults to False.

Returns:

With a pipe, the received tensordict, which is self only if the tensordict is consolidated with the layout of the data, so use the return value. With consolidated=True, None. Otherwise, the tag of the last leaf received.

Raises:

ValueError – if both or none of src and group_src are given, or if group_src is given without group.

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