reference/generated/tensordict._ucxx.send_tensordict

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

send_tensordict

async tensordict._ucxx.send_tensordict(endpoint, td: TensorDictBase, *, consolidated: bool = True) → None

Sends a tensordict over a UCXX endpoint, with its metadata, in one call.

Warning

send_tensordict is experimental: its API can change in a minor release without a deprecation cycle.

Unlike a TensorDictPipe, it keeps no state: each call sends the metadata (as JSON) and then the consolidated storage. The receiver calls recv_tensordict(). It needs the optional ucxx package. The protocol does not authenticate the peer: use it only on a trusted network.

Parameters:

• endpoint (ucxx.Endpoint) – the endpoint to send through, for example from ucxx.create_endpoint.

• td (TensorDictBase) – the tensordict to send.

Keyword Arguments:

consolidated (bool, optional) – if True, consolidates the tensordict (see consolidate()) unless it is already, and sends its storage as one buffer. If False, sends the tensor leaves of the tensordict, nested ones included, one by one, in the format of a pipe with consolidated=False: the receiver calls recv_tensordict() with consolidated=False. Defaults to True.

Raises:

• ImportError – if ucxx is not installed.

• ValueError – with consolidated=False, if the tensordict holds a non-tensor entry other than a NonTensorData or NonTensorStack whose values are str, int, float, bool, None, or lists or dicts (with str keys) of them. Send it with consolidated=True.