LazyStackedTensorDict

class tensordict.LazyStackedTensorDict(*tensordicts: T, stack_dim: int = 0, hook_out: callable | None = None, hook_in: callable | None = None, batch_size: Sequence[int] | None = None, device: torch.device | None = None, names: Sequence[str] | None = None, stack_dim_name: str | None = None, strict_shape: bool = False)

TensorDict 的惰性堆叠。

当将 TensorDict 堆叠在一起时，默认行为是将它们放在一个未实例化的堆栈中。这允许无缝地处理 tensordict 堆栈，进行会影响原始 tensordict 的操作。

参数:

• *tensordicts (TensorDict 实例) – 具有相同批大小的 tensordict 列表。

• stack_dim (int) – 执行堆叠的维度（介于 -td.ndimension() 和 td.ndimension()-1 之间）。

• hook_out (callable, 可选) – 在 get() 之后执行的可调用对象。

• hook_in (callable, 可选) – 在 set() 之前执行的可调用对象。

• stack_dim_name – 堆叠维度的名称。默认为 None。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> tds = [TensorDict({'a': torch.randn(3, 4)}, batch_size=[3])
...     for _ in range(10)]
>>> td_stack = torch.stack(tds, -1)
>>> print(td_stack.shape)
torch.Size([3, 10])
>>> print(td_stack.get("a").shape)
torch.Size([3, 10, 4])
>>> print(td_stack[:, 0] is tds[0])
True

abs() → T

计算 TensorDict 中每个元素的绝对值。

abs_() → T

就地计算 TensorDict 中每个元素的绝对值。

acos() → T

计算 TensorDict 中每个元素的 acos() 值。

acos_() → T

就地计算 TensorDict 中每个元素的 acos() 值。

add(other: tensordict.base.TensorDictBase | torch.Tensor, *, alpha: Optional[float] = None, default: Optional[Union[str, Tensor]] = None) → TensorDictBase

将按 alpha 缩放的 other 加到 self。

\[\text{{out}}_i = \text{{input}}_i + \text{{alpha}} \times \text{{other}}_i\]

参数:

other (TensorDictBase 或 torch.Tensor) – 要加到 self 的张量或 TensorDict。

关键字参数:

• alpha (Number, 可选) – other 的乘数。

• default (torch.Tensor 或 str, 可选) – 用于独占条目的默认值。如果未提供，则两个 tensordict 的键列表必须完全匹配。如果传递了 default="intersection"，则仅考虑相交的键集，其他键将被忽略。在所有其他情况下，default 将用于操作两侧的所有缺失条目。

add_(other: tensordict.base.TensorDictBase | float, *, alpha: Optional[float] = None)

add() 的就地版本。

注意

就地 add 不支持 default 关键字参数。

addcdiv(other1: tensordict.base.TensorDictBase | torch.Tensor, other2: tensordict.base.TensorDictBase | torch.Tensor, value: float | None = 1)

执行 other1 除以 other2 的逐元素除法，将结果乘以标量 value 并将其加到 self。

\[\text{out}_i = \text{input}_i + \text{value} \times \frac{\text{tensor1}_i}{\text{tensor2}_i}\]

self、other1 和 other2 的元素的形状必须是可广播的。

对于 FloatTensor 或 DoubleTensor 类型的输入，value 必须是实数，否则为整数。

参数:

• other1 (TensorDict 或 Tensor) – 分子 tensordict（或张量）

• tensor2 (TensorDict 或 Tensor) – 分母 tensordict（或张量）

关键字参数:

value (Number, 可选) – \(\text{tensor1} / \text{tensor2}\) 的乘数

addcdiv_(other1, other2, *, value: float | None = 1)

addcdiv() 的就地版本。

addcmul(other1, other2, *, value: float | None = 1)

执行 other1 乘以 other2 的逐元素乘法，将结果乘以标量 value 并将其加到 self。

\[\text{out}_i = \text{input}_i + \text{value} \times \text{other1}_i \times \text{other2}_i\]

self、other1 和 other2 的形状必须是可广播的。

对于 FloatTensor 或 DoubleTensor 类型的输入，value 必须是实数，否则为整数。

参数:

• other1 (TensorDict 或 Tensor) – 要相乘的 tensordict 或张量

• other2 (TensorDict 或 Tensor) – 要相乘的 tensordict 或张量

关键字参数:

value (Number, 可选) – \(other1 .* other2\) 的乘数

addcmul_(other1, other2, *, value: float | None = 1)

addcmul() 的就地版本。

all(dim: Optional[int] = None) → bool | tensordict.base.TensorDictBase

检查 tensordict 中是否所有值都为 True/非空。

参数:

dim (int, 可选) – 如果为 None，则返回一个布尔值，指示是否所有张量都返回 tensor.all() == True。如果为整数，则仅当此维度与 tensordict 形状兼容时，才对指定的维度调用 all。

any(dim: Optional[int] = None) → bool | tensordict.base.TensorDictBase

检查 tensordict 中是否任何值都为 True/非空。

参数:

dim (int, 可选) – 如果为 None，则返回一个布尔值，指示是否所有张量都返回 tensor.any() == True。如果为整数，则仅当此维度与 tensordict 形状兼容时，才对指定的维度调用 all。

append(tensordict: T) → None

将 TensorDict 追加到堆栈上。

类似于 list.append。追加的 TensorDict 必须具有兼容的 batch_size 和设备。追加操作是就地操作，不返回任何内容。

参数:

tensordict (TensorDictBase) – 要追加到堆栈上的 TensorDict。

apply(fn: Callable, *others: T, batch_size: Optional[Sequence[int]] = None, device: torch.device | None = _NoDefault.ZERO, names: Optional[Sequence[str]] = _NoDefault.ZERO, inplace: bool = False, default: Any = _NoDefault.ZERO, filter_empty: Optional[bool] = None, propagate_lock: bool = False, call_on_nested: bool = False, out: Optional[TensorDictBase] = None, **constructor_kwargs) → Optional[T]

将可调用对象应用于 tensordict 中存储的所有值，并将它们设置在一个新的 tensordict 中。

可调用对象的签名必须是 Callable[Tuple[Tensor, ...], Optional[Union[Tensor, TensorDictBase]]]。

参数:

• fn (Callable) – 要应用于 tensordict 中张量的函数。

• *others (TensorDictBase 实例, 可选) – 如果提供，这些 tensordict 实例应具有与 self 结构匹配的结构。fn 参数应接收与 tensordict 数量（包括 self）一样多的未命名输入。如果其他 tensordict 缺少条目，则可以通过 default 关键字参数传递默认值。

关键字参数:

• batch_size (int 序列, 可选) – 如果提供，结果 TensorDict 将具有所需的 batch_size。batch_size 参数应与转换后的 batch_size 匹配。这是一个仅关键字参数。

• device (torch.device, 可选) – 结果设备，如果有。

• names (str 列表, 可选) – 新的维度名称，以防 batch_size 被修改。

• inplace (bool, 可选) – 如果为 True，则进行原地更改。默认为 False。这是一个仅关键字参数。

• default (Any, 可选) – 其他 tensordict 中缺少条目的默认值。如果未提供，缺少条目将引发 KeyError。

• filter_empty (bool, 可选) – 如果 True，则会滤除空的 tensordict。这也降低了计算成本，因为不会创建和销毁空数据结构。非张量数据被视为叶节点，因此即使函数未触及，也会保留在 tensordict 中。为了向后兼容，默认为 False。

• propagate_lock (bool, 可选) – 如果 True，则锁定的 tensordict 将生成另一个锁定的 tensordict。默认为 False。

• call_on_nested (bool, 可选) –

  如果 True，则将对第一级张量和容器（TensorDict 或 tensorclass）调用该函数。在这种情况下，func 负责将其调用传播到嵌套级别。这允许在将调用传播到嵌套 tensordict 时进行细粒度行为。如果 False，则仅在叶节点上调用该函数，并且 apply 将负责将函数分派到所有叶节点。

  >>> td = TensorDict({"a": {"b": [0.0, 1.0]}, "c": [1.0, 2.0]})
  >>> def mean_tensor_only(val):
  ...     if is_tensor_collection(val):
  ...         raise RuntimeError("Unexpected!")
  ...     return val.mean()
  >>> td_mean = td.apply(mean_tensor_only)
  >>> def mean_any(val):
  ...     if is_tensor_collection(val):
  ...         # Recurse
  ...         return val.apply(mean_any, call_on_nested=True)
  ...     return val.mean()
  >>> td_mean = td.apply(mean_any, call_on_nested=True)
  

• out (TensorDictBase, 可选) –

  用于写入结果的 tensordict。这可以用于避免创建新的 tensordict

  >>> td = TensorDict({"a": 0})
  >>> td.apply(lambda x: x+1, out=td)
  >>> assert (td==1).all()
  

  警告

  如果对 tensordict 执行的操作需要访问多个键才能进行单个计算，则提供等于 self 的 out 参数可能会导致操作提供静默的错误结果。例如

  >>> td = TensorDict({"a": 1, "b": 1})
  >>> td.apply(lambda x: x+td["a"])["b"] # Right!
  tensor(2)
  >>> td.apply(lambda x: x+td["a"], out=td)["b"] # Wrong!
  tensor(3)
  

• **constructor_kwargs – 要传递给 TensorDict 构造函数的其他关键字参数。

返回:

具有 transformed_in 张量的新 tensordict。

示例

>>> td = TensorDict({
...     "a": -torch.ones(3),
...     "b": {"c": torch.ones(3)}},
...     batch_size=[3])
>>> td_1 = td.apply(lambda x: x+1)
>>> assert (td_1["a"] == 0).all()
>>> assert (td_1["b", "c"] == 2).all()
>>> td_2 = td.apply(lambda x, y: x+y, td)
>>> assert (td_2["a"] == -2).all()
>>> assert (td_2["b", "c"] == 2).all()

注意

如果函数返回 None，则忽略该条目。这可以用于过滤 tensordict 中的数据

>>> td = TensorDict({"1": 1, "2": 2, "b": {"2": 2, "1": 1}}, [])
>>> def filter(tensor):
...     if tensor == 1:
...         return tensor
>>> td.apply(filter)
TensorDict(
    fields={
        1: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        b: TensorDict(
            fields={
                1: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

注意

无论输入类型如何，apply 方法都将返回 TensorDict 实例。要保持相同的类型，可以执行

>>> out = td.clone(False).update(td.apply(...))

apply_(fn: Callable, *others, **kwargs)

将可调用对象应用于 tensordict 中存储的所有值，并在原地重写它们。

参数:

• fn (Callable) – 要应用于 tensordict 中张量的函数。

• *others (TensorDictBase 序列, 可选) – 要使用的其他 tensordict。

关键字参数：请参阅 apply()。

返回:

应用该函数后的 self 或 self 的副本

asin() → T

计算 TensorDict 中每个元素的 asin() 值。

asin_() → T

原地计算 TensorDict 中每个元素的 asin() 值。

atan() → T

计算 TensorDict 中每个元素的 atan() 值。

atan_() → T

原地计算 TensorDict 中每个元素的 atan() 值。

auto_batch_size_(batch_dims: Optional[int] = None) → T

设置 tensordict 的最大批次大小，最多可达可选的 batch_dims。

参数:

batch_dims (int, 可选) – 如果提供，则批次大小最大为 batch_dims 长度。

返回:

self

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict({"a": torch.randn(3, 4, 5), "b": {"c": torch.randn(3, 4, 6)}}, batch_size=[])
>>> td.auto_batch_size_()
>>> print(td.batch_size)
torch.Size([3, 4])
>>> td.auto_batch_size_(batch_dims=1)
>>> print(td.batch_size)
torch.Size([3])

auto_device_() → T

如果设备是唯一的，则自动设置设备。

返回：具有编辑后的 device 属性的 self。

property batch_dims: int

tensordict 批次大小的长度。

返回:

描述 tensordict 维度数量的整数。

property batch_size: Size

TensorDict 的形状（或批次大小）。

tensordict 的形状对应于其包含的张量的公共前 N 个维度，其中 N 是任意数字。批次大小与“特征大小”形成对比，后者表示张量的语义相关形状。例如，一批视频可能具有形状 [B, T, C, W, H]，其中 [B, T] 是批次大小（批次和时间维度），而 [C, W, H] 是特征维度（通道和空间维度）。

TensorDict 的形状在初始化时由用户控制（即，它不是从张量形状推断出来的）。

如果新大小与 TensorDict 内容兼容，则可以动态编辑 batch_size。例如，始终允许将批次大小设置为空值。

返回:

描述 TensorDict 批次大小的 Size 对象。

示例

>>> data = TensorDict({
...     "key 0": torch.randn(3, 4),
...     "key 1": torch.randn(3, 5),
...     "nested": TensorDict({"key 0": torch.randn(3, 4)}, batch_size=[3, 4])},
...     batch_size=[3])
>>> data.batch_size = () # resets the batch-size to an empty value

bfloat16()

将所有张量转换为 torch.bfloat16。

bool()

将所有张量转换为 torch.bool。

bytes(*, count_duplicates: bool = True) → int

计算包含张量的字节数。

关键字参数:

count_duplicates (bool) – 是否将重复张量计为独立张量。如果为 False，则仅会丢弃严格相同的张量（来自公共基础张量的相同视图但具有不同 id 的张量将被计数两次）。默认为 True（每个张量都假定为单个副本）。

classmethod cat(input, dim=0, *, out=None)

沿给定维度将 tensordict 连接成单个 tensordict。

此调用等效于调用 torch.cat()，但与 torch.compile 兼容。

cat_from_tensordict(dim: int = 0, *, sorted: Optional[Union[bool, List[NestedKey]]] = None, out: Optional[Tensor] = None) → Tensor

将 tensordict 的所有条目连接到单个张量中。

参数:

dim (int, 可选) – 条目应沿其连接的维度。

关键字参数:

• sorted (bool 或 NestedKey 列表) – 如果为 True，则条目将按字母顺序连接。如果为 False（默认），将使用字典顺序。或者，可以提供键名称列表，并且张量将相应地连接。这会产生一些开销，因为将对照 tensordict 中的叶名称列表检查键列表。

• out (torch.Tensor, 可选) – cat 操作的可选目标张量。

cat_tensors(*keys: NestedKey, out_key: NestedKey, dim: int = 0, keep_entries: bool = False) → T

将条目连接到新条目中，并可能删除原始值。

参数:

keys (NestedKey 序列) – 要连接的条目。

关键字参数

out_key (NestedKey): 连接后的输入的新键名称。 keep_entries (bool, optional): 如果为 False，则将删除 keys 中的条目。

默认为 False。

dim (int, optional): 必须发生连接的维度。

默认为 0。

返回：self

示例

>>> td = TensorDict(a=torch.zeros(1), b=torch.ones(1))
>>> td.cat_tensors("a", "b", out_key="c")
>>> assert "a" not in td
>>> assert (td["c"] == torch.tensor([0, 1])).all()

ceil() → T

计算 TensorDict 中每个元素的 ceil() 值。

ceil_() → T

原地计算 TensorDict 中每个元素的 ceil() 值。

chunk(chunks: int, dim: int = 0) → tuple[tensordict.base.TensorDictBase, ...]

如果可能，将 tensordict 拆分为指定数量的块。

每个块都是输入 tensordict 的视图。

参数:

• chunks (int) – 要返回的块数

• dim (int, 可选) – 拆分 tensordict 的维度。默认为 0。

示例

>>> td = TensorDict({
...     'x': torch.arange(24).reshape(3, 4, 2),
... }, batch_size=[3, 4])
>>> td0, td1 = td.chunk(dim=-1, chunks=2)
>>> td0['x']
tensor([[[ 0,  1],
         [ 2,  3]],
        [[ 8,  9],
         [10, 11]],
        [[16, 17],
         [18, 19]]])

clamp_max(other: tensordict.base.TensorDictBase | torch.Tensor, *, default: Optional[Union[str, Tensor]] = None) → T

如果 self 的元素大于该值，则将其钳制为 other。

参数:

other (TensorDict 或 Tensor) – 其他输入 tensordict 或张量。

关键字参数:

default (torch.Tensor 或 str, 可选) – 用于独占条目的默认值。如果未提供，则两个 tensordict 的键列表必须完全匹配。如果传递了 default="intersection"，则仅考虑相交的键集，其他键将被忽略。在所有其他情况下，default 将用于操作两侧的所有缺失条目。

clamp_max_(other: tensordict.base.TensorDictBase | torch.Tensor) → T

clamp_max() 的原地版本。

注意

原地 clamp_max 不支持 default 关键字参数。

clamp_min(other: tensordict.base.TensorDictBase | torch.Tensor, default: Optional[Union[str, Tensor]] = None) → T

如果 self 的元素小于该值，则将其钳制为 other。

参数:

other (TensorDict 或 Tensor) – 其他输入 tensordict 或张量。

关键字参数:

default (torch.Tensor 或 str, 可选) – 用于独占条目的默认值。如果未提供，则两个 tensordict 的键列表必须完全匹配。如果传递了 default="intersection"，则仅考虑相交的键集，其他键将被忽略。在所有其他情况下，default 将用于操作两侧的所有缺失条目。

clamp_min_(other: tensordict.base.TensorDictBase | torch.Tensor) → T

注意

Inplace clamp_min does not support default keyword argument.

clear() → T

Erases the content of the tensordict.

clear_device_() → T

Clears the device of the tensordict.

返回：self

clone(recurse: bool = True, **kwargs) → T

Clones a TensorDictBase subclass instance onto a new TensorDictBase subclass of the same type.

To create a TensorDict instance from any other TensorDictBase subtype, call the to_tensordict() method instead.

参数:

recurse (bool, optional) – if True, each tensor contained in the TensorDict will be copied too. Otherwise only the TensorDict tree structure will be copied. Defaults to True.

注意

Unlike many other ops (pointwise arithmetic, shape operations, …) clone does not inherit the original lock attribute. This design choice is made such that a clone can be created to be modified, which is the most frequent usage.

complex128()

Casts all tensors to torch.complex128.

complex32()

Casts all tensors to torch.complex32.

complex64()

Casts all tensors to torch.complex64.

consolidate(filename: Optional[Union[Path, str]] = None, *, num_threads=0, device: Optional[device] = None, non_blocking: bool = False, inplace: bool = False, return_early: bool = False, use_buffer: bool = False, share_memory: bool = False, pin_memory: bool = False, metadata: bool = False) → None

Consolidates the tensordict content in a single storage for fast serialization.

参数:

filename (Path, optional) – an optional file path for a memory-mapped tensor to use as a storage for the tensordict.

关键字参数:

• num_threads (integer, optional) – the number of threads to use for populating the storage.

• device (torch.device, optional) – an optional device where the storage must be instantiated.

• non_blocking (bool, optional) – non_blocking argument passed to copy_().

• inplace (bool, optional) – if True, the resulting tensordict is the same as self with updated values. Defaults to False.

• return_early (bool, optional) – if True and num_threads>0, the method will return a future of the tensordict. The resulting tensordict can be queried using future.result().

• use_buffer (bool, optional) – if True and a filename is passed, an intermediate local buffer will be created in shared memory, and the data will be copied at the storage location as a last step. This may be faster than writing directly to a distant physical memory (e.g., NFS). Defaults to False.

• share_memory (bool, optional) – if True, the storage will be placed in shared memory. Defaults to False.

• pin_memory (bool, optional) – whether the consolidated data should be placed in pinned memory. Defaults to False.

• metadata (bool, optional) – if True, the metadata will be stored alongisde the common storage. If a filename is provided, this is without effect. Storing the metadata can be useful when one wants to control how serialization is achieved, as TensorDict handles the pickling/unpickling of consolidated TDs differently if the metadata is or isn’t available.

注意

If the tensordict is already consolidated, all arguments are ignored and self is returned. Call contiguous() to re-consolidate.

示例

>>> import pickle
>>> import tempfile
>>> import torch
>>> import tqdm
>>> from torch.utils.benchmark import Timer
>>> from tensordict import TensorDict
>>> data = TensorDict({"a": torch.zeros(()), "b": {"c": torch.zeros(())}})
>>> data_consolidated = data.consolidate()
>>> # check that the data has a single data_ptr()
>>> assert torch.tensor([
...     v.untyped_storage().data_ptr() for v in data_c.values(True, True)
... ]).unique().numel() == 1
>>> # Serializing the tensordict will be faster with data_consolidated
>>> with open("data.pickle", "wb") as f:
...    print("regular", Timer("pickle.dump(data, f)", globals=globals()).adaptive_autorange())
>>> with open("data_c.pickle", "wb") as f:
...     print("consolidated", Timer("pickle.dump(data_consolidated, f)", globals=globals()).adaptive_autorange())

contiguous() → T

Returns a new tensordict of the same type with contiguous values (or self if values are already contiguous).

copy()

Return a shallow copy of the tensordict (ie, copies the structure but not the data).

Equivalent to TensorDictBase.clone(recurse=False)

copy_(tensordict: T, non_blocking: bool = False) → T

See TensorDictBase.update_.

The non-blocking argument will be ignored and is just present for compatibility with torch.Tensor.copy_().

copy_at_(tensordict: T, idx: Union[None, int, slice, str, Tensor, List[Any], Tuple[Any, ...]], non_blocking: bool = False) → T

See TensorDictBase.update_at_.

cos() → T

Computes the cos() value of each element of the TensorDict.

cos_() → T

Computes the cos() value of each element of the TensorDict in-place.

cosh() → T

Computes the cosh() value of each element of the TensorDict.

cosh_() → T

Computes the cosh() value of each element of the TensorDict in-place.

cpu(**kwargs) → T

Casts a tensordict to CPU.

This function also supports all the keyword arguments of to().

create_nested(key)

Creates a nested tensordict of the same shape, device and dim names as the current tensordict.

If the value already exists, it will be overwritten by this operation. This operation is blocked in locked tensordicts.

示例

>>> data = TensorDict({}, [3, 4, 5])
>>> data.create_nested("root")
>>> data.create_nested(("some", "nested", "value"))
>>> print(data)
TensorDict(
    fields={
        root: TensorDict(
            fields={
            },
            batch_size=torch.Size([3, 4, 5]),
            device=None,
            is_shared=False),
        some: TensorDict(
            fields={
                nested: TensorDict(
                    fields={
                        value: TensorDict(
                            fields={
                            },
                            batch_size=torch.Size([3, 4, 5]),
                            device=None,
                            is_shared=False)},
                    batch_size=torch.Size([3, 4, 5]),
                    device=None,
                    is_shared=False)},
            batch_size=torch.Size([3, 4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3, 4, 5]),
    device=None,
    is_shared=False)

cuda(device: Optional[int] = None, **kwargs) → T

Casts a tensordict to a cuda device (if not already on it).

参数:

device (int, optional) – if provided, the cuda device on which the tensor should be cast.

This function also supports all the keyword arguments of to().

property data

Returns a tensordict containing the .data attributes of the leaf tensors.

data_ptr(*, storage: bool = False)

Returns the data_ptr of the tensordict leaves.

This can be useful to check if two tensordicts share the same data_ptr().

关键字参数:

storage (bool, optional) – if True, tensor.untyped_storage().data_ptr() will be called instead. Defaults to False.

示例

>>> from tensordict import TensorDict
>>> td = TensorDict(a=torch.randn(2), b=torch.randn(2), batch_size=[2])
>>> assert (td0.data_ptr() == td.data_ptr()).all()

注意

LazyStackedTensorDict instances will be displayed as nested tensordicts to reflect the true data_ptr() of their leaves

>>> td0 = TensorDict(a=torch.randn(2), b=torch.randn(2), batch_size=[2])
>>> td1 = TensorDict(a=torch.randn(2), b=torch.randn(2), batch_size=[2])
>>> td = TensorDict.lazy_stack([td0, td1])
>>> td.data_ptr()
TensorDict(
    fields={
        0: TensorDict(
            fields={
                a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
                b: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
            batch_size=torch.Size([]),
            device=cpu,
            is_shared=False),
        1: TensorDict(
            fields={
                a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
                b: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
            batch_size=torch.Size([]),
            device=cpu,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=cpu,
    is_shared=False)

del_(key: NestedKey, **kwargs: Any) → T

Deletes a key of the tensordict.

参数:

key (NestedKey) – key to be deleted

返回:

self

densify(*, layout: layout = torch.strided)

Attempts to represent the lazy stack with contiguous tensors (plain tensors or nested).

关键字参数:

layout (torch.layout) – the layout of the nested tensors, if any. Defaults to strided.

property depth: int

Returns the depth - maximum number of levels - of a tensordict.

The minimum depth is 0 (no nested tensordict).

detach() → T

Detach the tensors in the tensordict.

返回:

a new tensordict with no tensor requiring gradient.

detach_() → T

Detach the tensors in the tensordict in-place.

返回:

self.

property device: torch.device | None

Device of a TensorDict.

If the TensorDict has a specified device, all its tensors (incl. nested ones) must live on the same device. If the TensorDict device is None, different values can be located on different devices.

返回:

torch.device object indicating the device where the tensors are placed, or None if TensorDict does not have a device.

示例

>>> td = TensorDict({
...     "cpu": torch.randn(3, device='cpu'),
...     "cuda": torch.randn(3, device='cuda'),
... }, batch_size=[], device=None)
>>> td['cpu'].device
device(type='cpu')
>>> td['cuda'].device
device(type='cuda')
>>> td = TensorDict({
...     "x": torch.randn(3, device='cpu'),
...     "y": torch.randn(3, device='cuda'),
... }, batch_size=[], device='cuda')
>>> td['x'].device
device(type='cuda')
>>> td['y'].device
device(type='cuda')
>>> td = TensorDict({
...     "x": torch.randn(3, device='cpu'),
...     "y": TensorDict({'z': torch.randn(3, device='cpu')}, batch_size=[], device=None),
... }, batch_size=[], device='cuda')
>>> td['x'].device
device(type='cuda')
>>> td['y'].device # nested tensordicts are also mapped onto the appropriate device.
device(type='cuda')
>>> td['y', 'x'].device
device(type='cuda')

dim() → int

See batch_dims().

div(other: tensordict.base.TensorDictBase | torch.Tensor, *, default: Optional[Union[str, Tensor]] = None) → T

Divides each element of the input self by the corresponding element of other.

\[\text{out}_i = \frac{\text{input}_i}{\text{other}_i}\]

Supports broadcasting, type promotion and integer, float, tensordict or tensor inputs. Always promotes integer types to the default scalar type.

参数:

other (TensorDict, Tensor or Number) – the divisor.

关键字参数:

default (torch.Tensor 或 str, 可选) – 用于独占条目的默认值。如果未提供，则两个 tensordict 的键列表必须完全匹配。如果传递了 default="intersection"，则仅考虑相交的键集，其他键将被忽略。在所有其他情况下，default 将用于操作两侧的所有缺失条目。

div_(other: tensordict.base.TensorDictBase | torch.Tensor) → T

In-place version of div().

注意

Inplace div does not support default keyword argument.

double()

将所有张量转换为 torch.bool。

property dtype

Returns the dtype of the values in the tensordict, if it is unique.

dumps(prefix: Optional[str] = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False) → T

Saves the tensordict to disk.

This function is a proxy to memmap().

empty(recurse=False, *, batch_size=None, device=_NoDefault.ZERO, names=None) → T

返回一个新的、空的 tensordict，具有相同的设备和批次大小。

参数:

recurse (bool, 可选) – 如果 True，则将完整复制 TensorDict 的结构，但不包含内容。否则，仅复制根结构。默认为 False。

关键字参数:

• batch_size (torch.Size, 可选) – tensordict 的新批次大小。

• device (torch.device, 可选) – 新设备。

• names (list of str, 可选) – 维度名称。

entry_class(key: NestedKey) → type

返回条目的类，可能避免调用 isinstance(td.get(key), type)。

当 get() 的执行开销很大时，应优先使用此方法而不是 tensordict.get(key).shape。

erf() → T

计算 TensorDict 中每个元素的 erf() 值。

erf_() → T

就地计算 TensorDict 中每个元素的 erf() 值。

erfc() → T

计算 TensorDict 中每个元素的 erfc() 值。

erfc_() → T

就地计算 TensorDict 中每个元素的 erfc() 值。

exclude(*keys: NestedKey, inplace: bool = False) → T

排除 tensordict 的键，并返回不包含这些条目的新 tensordict。

值不会被复制：对原始或新 tensordict 的张量进行就地修改将导致两个 tensordict 都发生更改。

参数:

• *keys (str) – 要排除的键。

• inplace (bool) – 如果为 True，则就地修剪 tensordict。默认为 False。

返回:

一个新的 tensordict（如果 inplace=True 则为同一个）不包含排除的条目。

示例

>>> from tensordict import TensorDict
>>> td = TensorDict({"a": 0, "b": {"c": 1, "d": 2}}, [])
>>> td.exclude("a", ("b", "c"))
TensorDict(
    fields={
        b: TensorDict(
            fields={
                d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> td.exclude("a", "b")
TensorDict(
    fields={
    },
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

exp() → T

计算 TensorDict 中每个元素的 exp() 值。

exp_() → T

就地计算 TensorDict 中每个元素的 exp() 值。

expand(*args: int, inplace: bool = False) → T

根据 expand() 函数扩展 tensordict 的每个张量，忽略特征维度。

支持使用可迭代对象来指定形状。

示例

>>> td = TensorDict({
...     'a': torch.zeros(3, 4, 5),
...     'b': torch.zeros(3, 4, 10)}, batch_size=[3, 4])
>>> td_expand = td.expand(10, 3, 4)
>>> assert td_expand.shape == torch.Size([10, 3, 4])
>>> assert td_expand.get("a").shape == torch.Size([10, 3, 4, 5])

expand_as(other: tensordict.base.TensorDictBase | torch.Tensor) → TensorDictBase

将 tensordict 的形状广播到 other 的形状，并据此进行扩展。

如果输入是张量集合（tensordict 或 tensorclass），则叶节点将逐一扩展。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td0 = TensorDict({
...     "a": torch.ones(3, 1, 4),
...     "b": {"c": torch.ones(3, 2, 1, 4)}},
...     batch_size=[3],
... )
>>> td1 = TensorDict({
...     "a": torch.zeros(2, 3, 5, 4),
...     "b": {"c": torch.zeros(2, 3, 2, 6, 4)}},
...     batch_size=[2, 3],
... )
>>> expanded = td0.expand_as(td1)
>>> assert (expanded==1).all()
>>> print(expanded)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([2, 3, 5, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([2, 3, 2, 6, 4]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([2, 3]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([2, 3]),
    device=None,
    is_shared=False)

expm1() → T

计算 TensorDict 中每个元素的 expm1() 值。

expm1_() → T

就地计算 TensorDict 中每个元素的 expm1() 值。

fill_(key: NestedKey, value: float | bool) → T

用给定的标量值填充键指向的张量。

参数:

• key (str 或 nested key) – 要填充的条目。

• value (Number 或 bool) – 用于填充的值。

返回:

self

filter_empty_()

就地过滤掉所有空的 tensordict。

filter_non_tensor_data() → T

过滤掉所有非张量数据。

flatten(start_dim=0, end_dim=- 1)

展平 tensordict 的所有张量。

参数:

• start_dim (int) – 要展平的第一个维度

• end_dim (int) – 要展平的最后一个维度

示例

>>> td = TensorDict({
...     "a": torch.arange(60).view(3, 4, 5),
...     "b": torch.arange(12).view(3, 4)}, batch_size=[3, 4])
>>> td_flat = td.flatten(0, 1)
>>> td_flat.batch_size
torch.Size([12])
>>> td_flat["a"]
tensor([[ 0,  1,  2,  3,  4],
        [ 5,  6,  7,  8,  9],
        [10, 11, 12, 13, 14],
        [15, 16, 17, 18, 19],
        [20, 21, 22, 23, 24],
        [25, 26, 27, 28, 29],
        [30, 31, 32, 33, 34],
        [35, 36, 37, 38, 39],
        [40, 41, 42, 43, 44],
        [45, 46, 47, 48, 49],
        [50, 51, 52, 53, 54],
        [55, 56, 57, 58, 59]])
>>> td_flat["b"]
tensor([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11])

flatten_keys(separator: str = '.', inplace: bool = False, is_leaf: Optional[Callable[[Type], bool]] = None) → T

递归地将嵌套的 tensordict 转换为扁平的 tensordict。

TensorDict 类型将丢失，结果将是一个简单的 TensorDict 实例。

参数:

• separator (str, 可选) – 嵌套项之间的分隔符。

• inplace (bool, 可选) – 如果 True，则生成的 tensordict 将与进行调用的 tensordict 具有相同的标识。默认为 False。

• is_leaf (callable, 可选) – 一个可调用对象，针对类类型返回一个布尔值，指示该类是否必须被视为叶节点。

示例

>>> data = TensorDict({"a": 1, ("b", "c"): 2, ("e", "f", "g"): 3}, batch_size=[])
>>> data.flatten_keys(separator=" - ")
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        b - c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        e - f - g: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

此方法和 unflatten_keys() 在处理状态字典时特别有用，因为它们可以无缝地将扁平字典转换为模仿模型结构的数据结构。

示例

>>> model = torch.nn.Sequential(torch.nn.Linear(3 ,4))
>>> ddp_model = torch.ao.quantization.QuantWrapper(model)
>>> state_dict = TensorDict(ddp_model.state_dict(), batch_size=[]).unflatten_keys(".")
>>> print(state_dict)
TensorDict(
    fields={
        module: TensorDict(
            fields={
                0: TensorDict(
                    fields={
                        bias: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False),
                        weight: Tensor(shape=torch.Size([4, 3]), device=cpu, dtype=torch.float32, is_shared=False)},
                    batch_size=torch.Size([]),
                    device=None,
                    is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> model_state_dict = state_dict.get("module")
>>> print(model_state_dict)
TensorDict(
    fields={
        0: TensorDict(
            fields={
                bias: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False),
                weight: Tensor(shape=torch.Size([4, 3]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> model.load_state_dict(dict(model_state_dict.flatten_keys(".")))

float()

将所有张量转换为 torch.float。

float16()

将所有张量转换为 torch.float16。

float32()

将所有张量转换为 torch.float32。

float64()

将所有张量转换为 torch.float64。

floor() → T

计算 TensorDict 中每个元素的 floor() 值。

floor_() → T

就地计算 TensorDict 中每个元素的 floor() 值。

frac() → T

计算 TensorDict 中每个元素的 frac() 值。

frac_() → T

就地计算 TensorDict 中每个元素的 frac() 值。

classmethod from_dict(input_dict, batch_size=None, device=None, batch_dims=None, stack_dim_name=None, stack_dim=0)

返回从字典或另一个 TensorDict 创建的 TensorDict。

如果未指定 batch_size，则返回可能的最大批次大小。

此函数也适用于嵌套字典，或者可用于确定嵌套 tensordict 的批次大小。

参数:

• input_dict (dictionary, 可选) – 用作数据源的字典（兼容嵌套键）。

• batch_size (iterable of int, 可选) – tensordict 的批次大小。

• device (torch.device 或 兼容类型, 可选) – TensorDict 的设备。

• batch_dims (int, 可选) – batch_dims（即要考虑用于 batch_size 的前导维度数）。与 batch_size 互斥。请注意，这是 tensordict 的 __最大__ 批次维度数，允许较小的数字。

• names (list of str, 可选) – tensordict 的维度名称。

示例

>>> input_dict = {"a": torch.randn(3, 4), "b": torch.randn(3)}
>>> print(TensorDict.from_dict(input_dict))
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        b: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([3]),
    device=None,
    is_shared=False)
>>> # nested dict: the nested TensorDict can have a different batch-size
>>> # as long as its leading dims match.
>>> input_dict = {"a": torch.randn(3), "b": {"c": torch.randn(3, 4)}}
>>> print(TensorDict.from_dict(input_dict))
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 4]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3]),
    device=None,
    is_shared=False)
>>> # we can also use this to work out the batch sie of a tensordict
>>> input_td = TensorDict({"a": torch.randn(3), "b": {"c": torch.randn(3, 4)}}, [])
>>> print(TensorDict.from_dict(input_td))
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 4]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3]),
    device=None,
    is_shared=False)

from_dict_instance(input_dict, batch_size=None, device=None, batch_dims=None, names=None)

from_dict() 的实例方法版本。

与 from_dict() 不同，此方法将尝试在现有树中保留 tensordict 类型（对于任何现有叶节点）。

示例

>>> from tensordict import TensorDict, tensorclass
>>> import torch
>>>
>>> @tensorclass
>>> class MyClass:
...     x: torch.Tensor
...     y: int
>>>
>>> td = TensorDict({"a": torch.randn(()), "b": MyClass(x=torch.zeros(()), y=1)})
>>> print(td.from_dict_instance(td.to_dict()))
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: MyClass(
            x=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
            y=Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> print(td.from_dict(td.to_dict()))
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                x: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                y: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

classmethod from_h5(filename, mode='r')

从 h5 文件创建 PersistentTensorDict。

此函数将自动确定每个嵌套 tensordict 的批次大小。

参数:

• filename (str) – h5 文件的路径。

• mode (str, 可选) – 读取模式。默认为 "r"。

classmethod from_module(module, as_module: bool = False, lock: bool = True, use_state_dict: bool = False)

将模块的参数和缓冲区复制到 tensordict 中。

参数:

• module (nn.Module) – 要从中获取参数的模块。

• as_module (bool, 可选) – 如果 True，将返回一个 TensorDictParams 实例，该实例可用于在 torch.nn.Module 中存储参数。默认为 False。

• lock (bool, 可选) – 如果 True，则生成的 tensordict 将被锁定。默认为 True。

• use_state_dict (bool, 可选) –

  如果 True，将使用模块的状态字典并将其解展平为具有模型树结构的 TensorDict。默认为 False。

  注意

  当必须使用状态字典钩子时，这尤其有用。

示例

>>> from torch import nn
>>> module = nn.TransformerDecoder(
...     decoder_layer=nn.TransformerDecoderLayer(nhead=4, d_model=4),
...     num_layers=1
... )
>>> params = TensorDict.from_module(module)
>>> print(params["layers", "0", "linear1"])
TensorDict(
    fields={
        bias: Parameter(shape=torch.Size([2048]), device=cpu, dtype=torch.float32, is_shared=False),
        weight: Parameter(shape=torch.Size([2048, 4]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

classmethod from_modules(*modules, as_module: bool = False, lock: bool = True, use_state_dict: bool = False, lazy_stack: bool = False, expand_identical: bool = False)

检索多个模块的参数，用于 ensebmle 学习/expects 应用的特征，通过 vmap。

参数:

modules (sequence of nn.Module) – 要从中获取参数的模块序列。如果模块的结构不同，则需要延迟堆叠（请参阅下面的 lazy_stack 参数）。

关键字参数:

• as_module (bool, 可选) – 如果 True，将返回一个 TensorDictParams 实例，该实例可用于在 torch.nn.Module 中存储参数。默认为 False。

• lock (bool, 可选) – 如果 True，则生成的 tensordict 将被锁定。默认为 True。

• use_state_dict (bool, 可选) –

  如果 True，将使用模块的状态字典并将其解展平为具有模型树结构的 TensorDict。默认为 False。

  注意

  当必须使用状态字典钩子时，这尤其有用。

• lazy_stack (bool, 可选) –

  参数应该是密集堆叠还是延迟堆叠。默认为 False（密集堆叠）。

  注意

  lazy_stack 和 as_module 是互斥的功能。

  警告

  延迟输出和非延迟输出之间存在关键区别，即非延迟输出将使用所需的批次大小重新实例化参数，而 lazy_stack 将仅将参数表示为延迟堆叠。这意味着，当 lazy_stack=True 时，原始参数可以安全地传递给优化器，而当设置为 True 时，则需要传递新参数。

  警告

  虽然使用延迟堆叠来保留原始参数引用可能很诱人，但请记住，延迟堆叠在每次调用 get() 时都会执行堆叠。这将需要内存（参数大小的 N 倍，如果构建图形则更多）和时间来计算。这也意味着优化器将包含更多参数，并且诸如 step() 或 zero_grad() 之类的操作将花费更长的时间才能执行。通常，lazy_stack 应保留给极少数用例。

• expand_identical (bool, 可选) – 如果 True 并且相同的参数（相同的标识）正在堆叠到自身，则将返回此参数的扩展版本。当 lazy_stack=True 时，此参数将被忽略。

示例

>>> from torch import nn
>>> from tensordict import TensorDict
>>> torch.manual_seed(0)
>>> empty_module = nn.Linear(3, 4, device="meta")
>>> n_models = 2
>>> modules = [nn.Linear(3, 4) for _ in range(n_models)]
>>> params = TensorDict.from_modules(*modules)
>>> print(params)
TensorDict(
    fields={
        bias: Parameter(shape=torch.Size([2, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        weight: Parameter(shape=torch.Size([2, 4, 3]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([2]),
    device=None,
    is_shared=False)
>>> # example of batch execution
>>> def exec_module(params, x):
...     with params.to_module(empty_module):
...         return empty_module(x)
>>> x = torch.randn(3)
>>> y = torch.vmap(exec_module, (0, None))(params, x)
>>> assert y.shape == (n_models, 4)
>>> # since lazy_stack = False, backprop leaves the original params untouched
>>> y.sum().backward()
>>> assert params["weight"].grad.norm() > 0
>>> assert modules[0].weight.grad is None

使用 lazy_stack=True 时，情况略有不同

>>> params = TensorDict.from_modules(*modules, lazy_stack=True)
>>> print(params)
LazyStackedTensorDict(
    fields={
        bias: Tensor(shape=torch.Size([2, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        weight: Tensor(shape=torch.Size([2, 4, 3]), device=cpu, dtype=torch.float32, is_shared=False)},
    exclusive_fields={
    },
    batch_size=torch.Size([2]),
    device=None,
    is_shared=False,
    stack_dim=0)
>>> # example of batch execution
>>> y = torch.vmap(exec_module, (0, None))(params, x)
>>> assert y.shape == (n_models, 4)
>>> y.sum().backward()
>>> assert modules[0].weight.grad is not None

classmethod from_namedtuple(named_tuple, *, auto_batch_size: bool = False)

递归地将 namedtuple 转换为 TensorDict。

关键字参数:

auto_batch_size (bool, 可选) – 如果 True，将自动计算批次大小。默认为 False。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> data = TensorDict({
...     "a_tensor": torch.zeros((3)),
...     "nested": {"a_tensor": torch.zeros((3)), "a_string": "zero!"}}, [3])
>>> nt = data.to_namedtuple()
>>> print(nt)
GenericDict(a_tensor=tensor([0., 0., 0.]), nested=GenericDict(a_tensor=tensor([0., 0., 0.]), a_string='zero!'))
>>> TensorDict.from_namedtuple(nt, auto_batch_size=True)
TensorDict(
    fields={
        a_tensor: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
        nested: TensorDict(
            fields={
                a_string: NonTensorData(data=zero!, batch_size=torch.Size([3]), device=None),
                a_tensor: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3]),
    device=None,
    is_shared=False)

classmethod from_pytree(pytree, *, batch_size: Optional[Size] = None, auto_batch_size: bool = False, batch_dims: Optional[int] = None)

将 pytree 转换为 TensorDict 实例。

此方法旨在尽可能保留 pytree 的嵌套结构。

添加额外的非张量键以跟踪每个级别的标识，从而提供内置的 pytree 到 tensordict 的双射转换 API。

目前接受的类包括列表、元组、命名元组和字典。

注意

对于字典，非 NestedKey 键将作为 NonTensorData 实例单独注册。

注意

可转换为张量的类型（例如 int、float 或 np.ndarray）将转换为 torch.Tensor 实例。请注意，此转换是满射的：将 tensordict 转换回 pytree 将不会恢复原始类型。

示例

>>> # Create a pytree with tensor leaves, and one "weird"-looking dict key
>>> class WeirdLookingClass:
...     pass
...
>>> weird_key = WeirdLookingClass()
>>> # Make a pytree with tuple, lists, dict and namedtuple
>>> pytree = (
...     [torch.randint(10, (3,)), torch.zeros(2)],
...     {
...         "tensor": torch.randn(
...             2,
...         ),
...         "td": TensorDict({"one": 1}),
...         weird_key: torch.randint(10, (2,)),
...         "list": [1, 2, 3],
...     },
...     {"named_tuple": TensorDict({"two": torch.ones(1) * 2}).to_namedtuple()},
... )
>>> # Build a TensorDict from that pytree
>>> td = TensorDict.from_pytree(pytree)
>>> # Recover the pytree
>>> pytree_recon = td.to_pytree()
>>> # Check that the leaves match
>>> def check(v1, v2):
>>>     assert (v1 == v2).all()
>>>
>>> torch.utils._pytree.tree_map(check, pytree, pytree_recon)
>>> assert weird_key in pytree_recon[1]

classmethod from_struct_array(struct_array: ndarray, device: Optional[device] = None) → T

将结构化的 numpy 数组转换为 TensorDict。

生成的 TensorDict 的内容将与 numpy 数组共享相同的内存内容（这是一个零拷贝操作）。就地更改结构化的 numpy 数组的值将影响 TensorDict 的内容。

示例

>>> x = np.array(
...     [("Rex", 9, 81.0), ("Fido", 3, 27.0)],
...     dtype=[("name", "U10"), ("age", "i4"), ("weight", "f4")],
... )
>>> td = TensorDict.from_struct_array(x)
>>> x_recon = td.to_struct_array()
>>> assert (x_recon == x).all()
>>> assert x_recon.shape == x.shape
>>> # Try modifying x age field and check effect on td
>>> x["age"] += 1
>>> assert (td["age"] == np.array([10, 4])).all()

classmethod fromkeys(keys: List[NestedKey], value: Any = 0)

从键列表和单个值创建 tensordict。

参数:

• keys (NestedKey 列表) – 一个可迭代对象，用于指定新字典的键。

• value (兼容类型, 可选) – 所有键的值。默认为 0。

gather(dim: int, index: Tensor, out: Optional[T] = None) → T

沿 dim 指定的轴收集值。

参数:

• dim (int) – 收集元素的维度

• index (torch.Tensor) – 一个长张量，其维度数量与 tensordict 的维度数量匹配，但只有一个维度不同（收集维度）。其元素指的是沿所需维度收集的索引。

• out (TensorDictBase, 可选) – 目标 tensordict。它必须具有与索引相同的形状。

示例

>>> td = TensorDict(
...     {"a": torch.randn(3, 4, 5),
...      "b": TensorDict({"c": torch.zeros(3, 4, 5)}, [3, 4, 5])},
...     [3, 4])
>>> index = torch.randint(4, (3, 2))
>>> td_gather = td.gather(dim=1, index=index)
>>> print(td_gather)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 2, 5]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([3, 2, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 2, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3, 2]),
    device=None,
    is_shared=False)

Gather 保留维度名称。

示例

>>> td.names = ["a", "b"]
>>> td_gather = td.gather(dim=1, index=index)
>>> td_gather.names
["a", "b"]

gather_and_stack(dst: int, group: 'dist.ProcessGroup' | None = None) → T | None

从不同的工作进程收集 tensordict，并将它们堆叠到目标工作进程的 self 上。

参数:

• dst (int) – 调用 gather_and_stack() 的目标工作进程的 rank。

• group (torch.distributed.ProcessGroup, 可选) – 如果设置，则指定的进程组将用于通信。否则，将使用默认进程组。默认为 None。

示例

>>> from torch import multiprocessing as mp
>>> from tensordict import TensorDict
>>> import torch
>>>
>>> def client():
...     torch.distributed.init_process_group(
...         "gloo",
...         rank=1,
...         world_size=2,
...         init_method=f"tcp://127.0.0.1:10003",
...     )
...     # Create a single tensordict to be sent to server
...     td = TensorDict(
...         {("a", "b"): torch.randn(2),
...          "c": torch.randn(2)}, [2]
...     )
...     td.gather_and_stack(0)
...
>>> def server():
...     torch.distributed.init_process_group(
...         "gloo",
...         rank=0,
...         world_size=2,
...         init_method=f"tcp://127.0.0.1:10003",
...     )
...     # Creates the destination tensordict on server.
...     # The first dim must be equal to world_size-1
...     td = TensorDict(
...         {("a", "b"): torch.zeros(2),
...          "c": torch.zeros(2)}, [2]
...     ).expand(1, 2).contiguous()
...     td.gather_and_stack(0)
...     assert td["a", "b"] != 0
...     print("yuppie")
...
>>> if __name__ == "__main__":
...     mp.set_start_method("spawn")
...
...     main_worker = mp.Process(target=server)
...     secondary_worker = mp.Process(target=client)
...
...     main_worker.start()
...     secondary_worker.start()
...
...     main_worker.join()
...     secondary_worker.join()

get(key: NestedKey, *args, **kwargs) → Tensor

获取使用输入键存储的值。

参数:

• key (str, str 元组) – 要查询的键。如果是 str 元组，则等效于 getattr 的链式调用。

• default –

  如果在 tensordict 中找不到键，则为默认值。

  警告

  目前，如果键在 tensordict 中不存在且未传递默认值，则会引发 KeyError。从 v0.7 开始，此行为将更改，并将改为返回 None 值。要采用新行为，请设置环境变量 export TD_GET_DEFAULTS_TO_NONE=’1’ 或调用 :func`~tensordict.set_get_defaults_to_none`。

示例

>>> td = TensorDict({"x": 1}, batch_size=[])
>>> td.get("x")
tensor(1)
>>> set_get_defaults_to_none(False) # Current default behaviour
>>> td.get("y") # Raises KeyError
>>> set_get_defaults_to_none(True)
>>> td.get("y")
None

get_at(key: NestedKey, index: Union[None, int, slice, str, Tensor, List[Any], Tuple[Any, ...]], default: Tensor = _NoDefault.ZERO) → Tensor

从键 key 在索引 idx 处获取 tensordict 的值。

参数:

• key (str, str 元组) – 要检索的键。

• index (int, slice, torch.Tensor, 可迭代对象) – 张量的索引。

• default (torch.Tensor) – 如果键在 tensordict 中不存在，则返回的默认值。

返回:

索引张量。

示例

>>> td = TensorDict({"x": torch.arange(3)}, batch_size=[])
>>> td.get_at("x", index=1)
tensor(1)

get_item_shape(key)

获取惰性堆栈中项目的形状。

异构维度返回为 -1。

此实现效率低下，因为它将尝试堆叠项目以计算其形状，仅应用于打印。

get_nestedtensor(key: NestedKey, default: Any = _NoDefault.ZERO, *, layout: Optional[layout] = None) → Tensor

当无法实现堆叠时，返回嵌套张量。

参数:

• key (NestedKey) – 要嵌套的条目。

• default (Any, 可选) –

  如果键不在所有子 tensordict 中，则返回的默认值。

  注意

  如果默认值是张量，则此方法将尝试使用它构造嵌套张量。否则，将返回默认值。

关键字参数:

layout (torch.layout, 可选) – 嵌套张量的布局。

示例

>>> td0 = TensorDict({"a": torch.zeros(4), "b": torch.zeros(4)}, [])
>>> td1 = TensorDict({"a": torch.ones(5)}, [])
>>> td = torch.stack([td0, td1], 0)
>>> a = td.get_nestedtensor("a")
>>> # using a tensor as default uses this default to build the nested tensor
>>> b = td.get_nestedtensor("b", default=torch.ones(4))
>>> assert (a == b).all()
>>> # using anything else as default returns the default
>>> b2 = td.get_nestedtensor("b", None)
>>> assert b2 is None

get_non_tensor(key: NestedKey, default=_NoDefault.ZERO)

获取非张量值（如果存在），如果找不到非张量值，则返回 default。

此方法对于张量/TensorDict 值是鲁棒的，这意味着如果收集的值是常规张量，它也会被返回（尽管此方法带有一些开销，不应在其自然范围之外使用）。

有关如何在 tensordict 中设置非张量值的更多信息，请参阅 set_non_tensor()。

参数:

• key (NestedKey) – NonTensorData 对象的位置。

• default (Any, 可选) – 如果找不到键，则要返回的值。

返回：tensordict.tensorclass.NonTensorData 的内容，

或者，如果条目不是 tensordict.tensorclass.NonTensorData，则返回与 key 对应的条目（如果找不到条目，则返回 default）。

示例

>>> data = TensorDict({}, batch_size=[])
>>> data.set_non_tensor(("nested", "the string"), "a string!")
>>> assert data.get_non_tensor(("nested", "the string")) == "a string!"
>>> # regular `get` works but returns a NonTensorData object
>>> data.get(("nested", "the string"))
NonTensorData(
    data='a string!',
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

property grad

返回一个 tensordict，其中包含叶张量的 .grad 属性。

half()

将所有张量转换为 torch.half。

insert(index: int, tensordict: T) → None

在指定索引处将 TensorDict 插入堆栈。

类似于 list.insert。插入的 TensorDict 必须具有兼容的 batch_size 和设备。插入是就地操作，不返回任何内容。

参数:

• index (int) – 应在其中插入新 TensorDict 的索引。

• tensordict (TensorDictBase) – 要插入堆栈的 TensorDict。

int()

将所有张量转换为 torch.int。

int16()

将所有张量转换为 torch.int16。

int32()

将所有张量转换为 torch.int32。

int64()

将所有张量转换为 torch.int64。

int8()

将所有张量转换为 torch.int8。

irecv(src: int, *, group: 'dist.ProcessGroup' | None = None, return_premature: bool = False, init_tag: int = 0, pseudo_rand: bool = False) → tuple[int, list[torch.Future]] | list[torch.Future] | None

异步接收 tensordict 的内容并使用其更新内容。

有关上下文，请查看 isend() 方法中的示例。

参数:

src (int) – 源工作进程的 rank。

关键字参数:

• group (torch.distributed.ProcessGroup, 可选) – 如果设置，则指定的进程组将用于通信。否则，将使用默认进程组。默认为 None。

• return_premature (bool) – 如果为 True，则返回一个 future 列表，用于等待直到 tensordict 更新。默认为 False，即等待直到在调用中完成更新。

• init_tag (int) – 源工作进程使用的 init_tag。

• pseudo_rand (bool) – 如果为 True，则标签序列将是伪随机的，允许从不同节点发送多个数据而不会重叠。请注意，这些伪随机数的生成成本很高（1e-5 秒/数字），这意味着它可能会减慢算法的运行时速度。此值必须与传递给 isend() 的值匹配。默认为 False。

返回:

如果 return_premature=True，则为 future 列表，用于等待

直到 tensordict 更新。

is_consolidated()

检查 TensorDict 是否具有合并存储。

is_contiguous() → bool

返回一个布尔值，指示所有张量是否是连续的。

is_empty() → bool

检查 tensordict 是否包含任何叶节点。

is_memmap() → bool

检查 tensordict 是否是内存映射的。

如果 TensorDict 实例是内存映射的，则它被锁定（条目不能被重命名、删除或添加）。如果使用都是内存映射的张量创建 TensorDict，这并__不__意味着 is_memmap 将返回 True（因为新张量可能是也可能不是内存映射的）。只有当调用 tensordict.memmap_() 时，tensordict 才会被视为内存映射的。

对于 CUDA 设备上的 tensordict，这始终为 True。

is_shared() → bool

检查 tensordict 是否在共享内存中。

如果 TensorDict 实例在共享内存中，则它被锁定（条目不能被重命名、删除或添加）。如果使用都是共享内存中的张量创建 TensorDict，这并__不__意味着 is_shared 将返回 True（因为新张量可能是也可能不在共享内存中）。只有当调用 tensordict.share_memory_() 或将 tensordict 放置在默认情况下内容共享的设备上（例如，"cuda"）时，tensordict 才会被视为在共享内存中。

对于 CUDA 设备上的 tensordict，这始终为 True。

isend(dst: int, *, group: 'dist.ProcessGroup' | None = None, init_tag: int = 0, pseudo_rand: bool = False) → int

异步发送 tensordict 的内容。

参数:

dst (int) – 目标工作进程的 rank，内容应发送到该进程。

关键字参数:

• group (torch.distributed.ProcessGroup, 可选) – 如果设置，则指定的进程组将用于通信。否则，将使用默认进程组。默认为 None。

• init_tag (int) – 用于标记张量的初始标签。请注意，这将递增 tensordict 中包含的张量数量。

• pseudo_rand (bool) – 如果为 True，则标签序列将是伪随机的，允许从不同节点发送多个数据而不会重叠。请注意，这些伪随机数的生成成本很高（1e-5 秒/数字），这意味着它可能会减慢算法的运行时速度。默认为 False。

示例

>>> import torch
>>> from tensordict import TensorDict
>>> from torch import multiprocessing as mp
>>> def client():
...     torch.distributed.init_process_group(
...         "gloo",
...         rank=1,
...         world_size=2,
...         init_method=f"tcp://127.0.0.1:10003",
...     )
...
...     td = TensorDict(
...         {
...             ("a", "b"): torch.randn(2),
...             "c": torch.randn(2, 3),
...             "_": torch.ones(2, 1, 5),
...         },
...         [2],
...     )
...     td.isend(0)
...
>>>
>>> def server(queue, return_premature=True):
...     torch.distributed.init_process_group(
...         "gloo",
...         rank=0,
...         world_size=2,
...         init_method=f"tcp://127.0.0.1:10003",
...     )
...     td = TensorDict(
...         {
...             ("a", "b"): torch.zeros(2),
...             "c": torch.zeros(2, 3),
...             "_": torch.zeros(2, 1, 5),
...         },
...         [2],
...     )
...     out = td.irecv(1, return_premature=return_premature)
...     if return_premature:
...         for fut in out:
...             fut.wait()
...     assert (td != 0).all()
...     queue.put("yuppie")
...
>>>
>>> if __name__ == "__main__":
...     queue = mp.Queue(1)
...     main_worker = mp.Process(
...         target=server,
...         args=(queue, )
...         )
...     secondary_worker = mp.Process(target=client)
...
...     main_worker.start()
...     secondary_worker.start()
...     out = queue.get(timeout=10)
...     assert out == "yuppie"
...     main_worker.join()
...     secondary_worker.join()

isfinite() → T

返回一个新的 tensordict，其中包含布尔元素，表示每个元素是否是有限的。

实数值在非 NaN、负无穷大或正无穷大时是有限的。复数值在其实部和虚部都有限时是有限的。

isnan() → T

返回一个新的 tensordict，其布尔元素表示输入的每个元素是否为 NaN。

当复数值的实部和/或虚部为 NaN 时，则认为复数值为 NaN。

isneginf() → T

测试输入的每个元素是否为负无穷大。

isposinf() → T

测试输入的每个元素是否为负无穷大。

isreal() → T

返回一个新的 tensordict，其布尔元素表示输入的每个元素是否为实数值。

items(include_nested=False, leaves_only=False, is_leaf=None, *, sort: bool = False)

返回 tensordict 的键值对生成器。

参数:

• include_nested (bool, optional) – 如果 True，将返回嵌套值。默认为 False。

• leaves_only (bool, optional) – 如果 False，则仅返回叶节点。默认为 False。

• is_leaf – 一个可选的可调用对象，指示一个类是否应被视为叶节点。

关键字参数:

sort (bool, optional) – 是否应排序键。对于嵌套键，键根据其连接名称排序（即，("a", "key") 将被视为 "a.key" 进行排序）。请注意，当处理大型 tensordict 时，排序可能会产生显著的开销。默认为 False。

keys(include_nested: bool = False, leaves_only: bool = False, is_leaf: Optional[Callable[[Type], bool]] = None, *, sort: bool = False) → _LazyStackedTensorDictKeysView

返回 tensordict 键的生成器。

警告

TensorDict 的 keys() 方法返回键的惰性视图。如果查询了 keys 但未对其进行迭代，然后修改了 tensordict，则稍后迭代键将返回键的新配置。

参数:

• include_nested (bool, optional) – 如果 True，将返回嵌套值。默认为 False。

• leaves_only (bool, optional) – 如果 False，则仅返回叶节点。默认为 False。

• is_leaf – 一个可选的可调用对象，指示一个类是否应被视为叶节点。

关键字参数:

sort (bool, optional) – 是否应对键进行排序。对于嵌套键，键根据其连接名称排序（即，("a", "key") 将被视为 "a.key" 进行排序）。请注意，当处理大型 tensordict 时，排序可能会产生显著的开销。默认为 False。

示例

>>> from tensordict import TensorDict
>>> data = TensorDict({"0": 0, "1": {"2": 2}}, batch_size=[])
>>> data.keys()
['0', '1']
>>> list(data.keys(leaves_only=True))
['0']
>>> list(data.keys(include_nested=True, leaves_only=True))
['0', '1', ('1', '2')]

classmethod lazy_stack(items: Sequence[TensorDictBase], dim: int = 0, *, device: Optional[Union[device, str, int]] = None, out: Optional[T] = None, stack_dim_name: Optional[str] = None, strict_shape: bool = False) → T

将 tensordict 堆叠到 LazyStackedTensorDict 中。

参数:

• items (TensorDictBase 实例序列) – 要堆叠的 TensorDictBase 实例序列。

• dim (int, optional) – 执行惰性堆叠的维度。默认为 0。

关键字参数:

• device (torch.device, optional) – 如果无法从 tensordict 列表推断设备（例如，列表为空），则在 LazyStackedTensorDict 中设置的设备。

• out (TensorDictBase, optional) – 用于写入数据的 LazyStackedTensorDict。

• stack_dim_name (str, optional) – 堆叠维度的名称。

• strict_shape (bool, optional) – 如果 True，则每个 tensordict 的形状都必须匹配。默认为 False。

lerp(end: tensordict.base.TensorDictBase | torch.Tensor, weight: tensordict.base.TensorDictBase | torch.Tensor | float)

基于标量或张量 weight，对两个张量 start（由 self 给定）和 end 进行线性插值。

\[\text{out}_i = \text{start}_i + \text{weight}_i \times (\text{end}_i - \text{start}_i)\]

start 和 end 的形状必须是可广播的。如果 weight 是张量，则 weight、start 和 end 的形状必须是可广播的。

参数:

• end (TensorDict) – 包含终点的 tensordict。

• weight (TensorDict, tensor 或 float) – 插值公式的权重。

lerp_(end: tensordict.base.TensorDictBase | torch.Tensor | float, weight: tensordict.base.TensorDictBase | torch.Tensor | float)

lerp() 的原地版本。

lgamma() → T

计算 TensorDict 中每个元素的 lgamma() 值。

lgamma_() → T

原地计算 TensorDict 中每个元素的 lgamma() 值。

classmethod load(prefix: str | pathlib.Path, *args, **kwargs) → T

从磁盘加载 tensordict。

此类方法是 load_memmap() 的代理。

load_(prefix: str | pathlib.Path, *args, **kwargs)

在当前 tensordict 中从磁盘加载 tensordict。

此类方法是 load_memmap_() 的代理。

classmethod load_memmap(prefix: str | pathlib.Path, device: Optional[device] = None, non_blocking: bool = False, *, out: Optional[TensorDictBase] = None) → T

从磁盘加载内存映射的 tensordict。

参数:

• prefix (str 或 文件夹路径) – 保存的 tensordict 应从中获取的文件夹路径。

• device (torch.device 或 等效项, optional) – 如果提供，数据将被异步转换为该设备。支持 “meta” 设备，在这种情况下，数据不会被加载，但会创建一组空的“meta”张量。这对于了解模型的总体大小和结构而无需实际打开任何文件非常有用。

• non_blocking (bool, optional) – 如果 True，则在设备上加载张量后不会调用 synchronize。默认为 False。

• out (TensorDictBase, optional) – 可选的 tensordict，数据应写入其中。

示例

>>> from tensordict import TensorDict
>>> td = TensorDict.fromkeys(["a", "b", "c", ("nested", "e")], 0)
>>> td.memmap("./saved_td")
>>> td_load = TensorDict.load_memmap("./saved_td")
>>> assert (td == td_load).all()

此方法还允许加载嵌套的 tensordict。

示例

>>> nested = TensorDict.load_memmap("./saved_td/nested")
>>> assert nested["e"] == 0

tensordict 也可以加载到 “meta” 设备上，或者作为伪张量加载。

示例

>>> import tempfile
>>> td = TensorDict({"a": torch.zeros(()), "b": {"c": torch.zeros(())}})
>>> with tempfile.TemporaryDirectory() as path:
...     td.save(path)
...     td_load = TensorDict.load_memmap(path, device="meta")
...     print("meta:", td_load)
...     from torch._subclasses import FakeTensorMode
...     with FakeTensorMode():
...         td_load = TensorDict.load_memmap(path)
...         print("fake:", td_load)
meta: TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=meta, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=meta, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([]),
            device=meta,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=meta,
    is_shared=False)
fake: TensorDict(
    fields={
        a: FakeTensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: FakeTensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([]),
            device=cpu,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=cpu,
    is_shared=False)

load_memmap_(prefix: str | pathlib.Path)

在调用 load_memmap_ 的 tensordict 中加载内存映射的 tensordict 的内容。

有关更多信息，请参阅 load_memmap()。

load_state_dict(state_dict: OrderedDict[str, Any], strict=True, assign=False, from_flatten=False) → T

将状态字典（格式如 state_dict() 中所示）加载到 tensordict 中。

参数:

• state_dict (OrderedDict) – 要复制的状态字典。

• strict (bool, optional) – 是否严格强制执行 state_dict 中的键与此 tensordict 的 torch.nn.Module.state_dict() 函数返回的键匹配。默认值：True

• assign (bool, optional) – 是否将状态字典中的项分配给 tensordict 中对应的键，而不是将它们原地复制到 tensordict 的当前张量中。当 False 时，保留当前模块中张量的属性，而当 True 时，保留状态字典中张量的属性。默认值：False

• from_flatten (bool, optional) – 如果 True，则假定输入 state_dict 是扁平化的。默认为 False。

示例

>>> data = TensorDict({"1": 1, "2": 2, "3": {"3": 3}}, [])
>>> data_zeroed = TensorDict({"1": 0, "2": 0, "3": {"3": 0}}, [])
>>> sd = data.state_dict()
>>> data_zeroed.load_state_dict(sd)
>>> print(data_zeroed["3", "3"])
tensor(3)
>>> # with flattening
>>> data_zeroed = TensorDict({"1": 0, "2": 0, "3": {"3": 0}}, [])
>>> data_zeroed.load_state_dict(data.state_dict(flatten=True), from_flatten=True)
>>> print(data_zeroed["3", "3"])
tensor(3)

lock_() → T

锁定 tensordict 以进行非原地操作。

诸如 set()、__setitem__()、update()、rename_key_() 或其他添加或删除条目的操作将被阻止。

此方法可以用作装饰器。

示例

>>> from tensordict import TensorDict
>>> td = TensorDict({"a": 1, "b": 2, "c": 3}, batch_size=[])
>>> with td.lock_():
...     assert td.is_locked
...     try:
...         td.set("d", 0) # error!
...     except RuntimeError:
...         print("td is locked!")
...     try:
...         del td["d"]
...     except RuntimeError:
...         print("td is locked!")
...     try:
...         td.rename_key_("a", "d")
...     except RuntimeError:
...         print("td is locked!")
...     td.set("a", 0, inplace=True)  # No storage is added, moved or removed
...     td.set_("a", 0) # No storage is added, moved or removed
...     td.update({"a": 0}, inplace=True)  # No storage is added, moved or removed
...     td.update_({"a": 0})  # No storage is added, moved or removed
>>> assert not td.is_locked

log() → T

计算 TensorDict 中每个元素的 log() 值。

log10() → T

计算 TensorDict 中每个元素的 log10() 值。

log10_() → T

原地计算 TensorDict 中每个元素的 log10() 值。

log1p() → T

计算 TensorDict 中每个元素的 log1p() 值。

log1p_() → T

原地计算 TensorDict 中每个元素的 log1p() 值。

log2() → T

计算 TensorDict 中每个元素的 log2() 值。

log2_() → T

原地计算 TensorDict 中每个元素的 log2() 值。

log_() → T

原地计算 TensorDict 中每个元素的 log() 值。

make_memmap(key: NestedKey, shape: torch.Size | torch.Tensor, *, dtype: Optional[dtype] = None) → MemoryMappedTensor

创建给定形状以及可能的 dtype 的空内存映射张量。

警告

此方法在设计上不是锁安全的。存在于多个节点上的内存映射 TensorDict 实例将需要使用方法 memmap_refresh_() 进行更新。

写入现有条目将导致错误。

参数:

• key (NestedKey) – 要写入的新条目的键。如果键已存在于 tensordict 中，则会引发异常。

• shape (torch.Size 或 等效项, 用于嵌套张量的 torch.Tensor) – 要写入的张量的形状。

关键字参数:

dtype (torch.dtype, 可选) – 新张量的数据类型。

返回:

一个新的内存映射张量。

make_memmap_from_storage(key: NestedKey, storage: UntypedStorage, shape: torch.Size | torch.Tensor, *, dtype: Optional[dtype] = None) → MemoryMappedTensor

创建给定存储、形状以及可能的 dtype 的空内存映射张量。

警告

此方法在设计上不是锁安全的。存在于多个节点上的内存映射 TensorDict 实例将需要使用方法 memmap_refresh_() 进行更新。

注意

如果存储具有关联的文件名，则它必须与该文件的新文件名匹配。如果它没有关联的文件名，但 tensordict 具有关联的路径，则这将导致异常。

参数:

• key (NestedKey) – 要写入的新条目的键。如果键已存在于 tensordict 中，则会引发异常。

• storage (torch.UntypedStorage) – 用于新 MemoryMappedTensor 的存储。必须是物理内存存储。

• shape (torch.Size 或 等效项, 用于嵌套张量的 torch.Tensor) – 要写入的张量的形状。

关键字参数:

dtype (torch.dtype, 可选) – 新张量的数据类型。

返回:

具有给定存储的新内存映射张量。

make_memmap_from_tensor(key: NestedKey, tensor: Tensor, *, copy_data: bool = True) → MemoryMappedTensor

创建给定张量的空内存映射张量。

警告

此方法在设计上不是锁安全的。存在于多个节点上的内存映射 TensorDict 实例将需要使用方法 memmap_refresh_() 进行更新。

如果 copy_data 为 True，则此方法始终复制存储内容（即，不共享存储）。

参数:

• key (NestedKey) – 要写入的新条目的键。如果键已存在于 tensordict 中，则会引发异常。

• tensor (torch.Tensor) – 要在物理内存上复制的张量。

关键字参数:

copy_data (bool, 可选) – 如果为 False，则新张量将共享输入的元数据（例如形状和 dtype），但内容将为空。默认为 True。

返回:

具有给定存储的新内存映射张量。

map(fn: Callable[[TensorDictBase], TensorDictBase | None], dim: int = 0, num_workers: int | None = None, *, out: TensorDictBase | None = None, chunksize: int | None = None, num_chunks: int | None = None, pool: mp.Pool | None = None, generator: torch.Generator | None = None, max_tasks_per_child: int | None = None, worker_threads: int = 1, index_with_generator: bool = False, pbar: bool = False, mp_start_method: str | None = None)

将函数映射到 tensordict 沿一个维度的分割。

此方法将通过将 tensordict 分块为大小相等的 tensordict，并将操作分派到所需数量的工作进程来将函数应用于 tensordict 实例。

函数签名应为 Callabe[[TensorDict], Union[TensorDict, Tensor]]。输出必须支持 torch.cat() 操作。该函数必须是可序列化的。

注意

当处理存储在磁盘上的大型数据集（例如，内存映射 tensordict）时，此方法特别有用，其中块将是原始数据的零拷贝切片，可以以几乎零成本传递给进程。 这允许以很小的成本处理非常大的数据集（例如，超过 Tb 大小）。

参数:

• fn (callable) – 要应用于 tensordict 的函数。支持类似于 Callabe[[TensorDict], Union[TensorDict, Tensor]] 的签名。

• dim (int, 可选) – tensordict 将沿其分块的维度。

• num_workers (int, 可选) – 工作进程的数量。与 pool 互斥。如果未提供，工作进程的数量将设置为可用 CPU 的数量。

关键字参数:

• out (TensorDictBase, 可选) – 输出的可选容器。其沿提供的 dim 的批量大小必须与 self.ndim 匹配。 如果它是共享的或 memmap 的（is_shared() 或 is_memmap() 返回 True），它将在远程进程中填充，从而避免数据向内传输。 否则，来自 self 切片的数据将被发送到进程，在当前进程上收集并就地写入 out。

• chunksize (int, 可选) – 每个数据块的大小。chunksize 为 0 将沿所需维度解绑 tensordict，并在应用该函数后重新堆叠它，而 chunksize>0 将拆分 tensordict 并在结果 tensordict 列表上调用 torch.cat()。 如果未提供，则块数将等于工作进程数。对于非常大的 tensordict，如此大的块可能无法容纳在内存中以完成操作，并且可能需要更多块才能使操作在实践中可行。此参数与 num_chunks 互斥。

• num_chunks (int, 可选) – 要将 tensordict 拆分成的块数。如果未提供，则块数将等于工作进程数。对于非常大的 tensordict，如此大的块可能无法容纳在内存中以完成操作，并且可能需要更多块才能使操作在实践中可行。此参数与 chunksize 互斥。

• pool (mp.Pool, 可选) – 用于执行作业的多进程池实例。如果未提供，将在 map 方法中创建池。

• generator (torch.Generator, 可选) –

  用于播种的生成器。将从中生成基本种子，并且池的每个工作进程将使用提供的种子（递增一个从 0 到 num_workers 的唯一整数）进行播种。 如果未提供生成器，则将使用随机整数作为种子。 要与未播种的工作进程一起工作，应单独创建池并将其直接传递给 map()。

  注意

  在提供低值种子时应谨慎，因为这可能会导致实验之间的自相关，例如：如果请求 8 个工作进程且种子为 4，则工作进程种子范围为 4 到 11。如果种子为 5，则工作进程种子范围为 5 到 12。 这两个实验将有 7 个种子的重叠，这可能会对结果产生意想不到的影响。

  注意

  为工作进程播种的目的是使每个工作进程都具有独立的种子，而不是在 map 方法的调用之间获得可重现的结果。换句话说，两个实验可能并且很可能会返回不同的结果，因为不可能知道哪个工作进程将选择哪个作业。 但是，我们可以确保每个工作进程都具有不同的种子，并且每个工作进程上的伪随机操作将是不相关的。

• max_tasks_per_child (int, 可选) – 每个子进程选取的最**大**作业数。默认为 None，即对作业数没有限制。

• worker_threads (int, 可选) – 工作进程的线程数。默认为 1。

• index_with_generator (bool, 可选) – 如果为 True，则 tensordict 的拆分/分块将在查询期间完成，从而节省初始化时间。 请注意，chunk() 和 split() 比索引（在生成器内使用）效率更高，因此在初始化时节省处理时间可能会对总运行时间产生负面影响。 默认为 False。

• pbar (bool, 可选) – 如果为 True，将显示进度条。需要安装 tqdm。 默认为 False。

• mp_start_method (str, 可选) – 多进程的启动方法。如果未提供，将使用默认启动方法。接受的字符串为 "fork" 和 "spawn"。 请记住，"cuda" 张量不能在使用 "fork" 启动方法的进程之间共享。如果 pool 传递给 map 方法，则这不起作用。

示例

>>> import torch
>>> from tensordict import TensorDict
>>>
>>> def process_data(data):
...     data.set("y", data.get("x") + 1)
...     return data
>>> if __name__ == "__main__":
...     data = TensorDict({"x": torch.zeros(1, 1_000_000)}, [1, 1_000_000]).memmap_()
...     data = data.map(process_data, dim=1)
...     print(data["y"][:, :10])
...
tensor([[1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]])

map_iter(fn: Callable[[TensorDictBase], TensorDictBase | None], dim: int = 0, num_workers: int | None = None, *, shuffle: bool = False, chunksize: int | None = None, num_chunks: int | None = None, pool: mp.Pool | None = None, generator: torch.Generator | None = None, max_tasks_per_child: int | None = None, worker_threads: int = 1, index_with_generator: bool = True, pbar: bool = False, mp_start_method: str | None = None)

迭代地将函数映射到 tensordict 沿一个维度的分割。

这是 map() 的可迭代版本。

此方法将通过将 tensordict 分块为大小相等的 tensordict，并将操作分派到所需数量的工作进程来将函数应用于 tensordict 实例。 它将一次产生一个结果。

函数签名应为 Callabe[[TensorDict], Union[TensorDict, Tensor]]。该函数必须是可序列化的。

注意

当处理存储在磁盘上的大型数据集（例如，内存映射 tensordict）时，此方法特别有用，其中块将是原始数据的零拷贝切片，可以以几乎零成本传递给进程。 这允许以很小的成本处理非常大的数据集（例如，超过 Tb 大小）。

注意

此函数用于表示数据集并从中加载数据，类似于数据加载器的方式。

参数:

• fn (callable) – 要应用于 tensordict 的函数。支持类似于 Callabe[[TensorDict], Union[TensorDict, Tensor]] 的签名。

• dim (int, 可选) – tensordict 将沿其分块的维度。

• num_workers (int, 可选) – 工作进程的数量。与 pool 互斥。如果未提供，工作进程的数量将设置为可用 CPU 的数量。

关键字参数:

• shuffle (bool, 可选) – 是否应全局打乱索引。如果为 True，则每个批次将包含非连续样本。如果 index_with_generator=False 且 shuffle=True`，则会引发错误。默认为 False。

• chunksize (int, 可选) – 每个数据块的大小。chunksize 为 0 将沿所需维度解绑 tensordict，并在应用该函数后重新堆叠它，而 chunksize>0 将拆分 tensordict 并在结果 tensordict 列表上调用 torch.cat()。 如果未提供，则块数将等于工作进程数。对于非常大的 tensordict，如此大的块可能无法容纳在内存中以完成操作，并且可能需要更多块才能使操作在实践中可行。此参数与 num_chunks 互斥。

• num_chunks (int, 可选) – 要将 tensordict 拆分成的块数。如果未提供，则块数将等于工作进程数。对于非常大的 tensordict，如此大的块可能无法容纳在内存中以完成操作，并且可能需要更多块才能使操作在实践中可行。此参数与 chunksize 互斥。

• pool (mp.Pool, 可选) – 用于执行作业的多进程池实例。如果未提供，将在 map 方法中创建池。

• generator (torch.Generator, 可选) –

  用于播种的生成器。将从中生成基本种子，并且池的每个工作进程将使用提供的种子（递增一个从 0 到 num_workers 的唯一整数）进行播种。 如果未提供生成器，则将使用随机整数作为种子。 要与未播种的工作进程一起工作，应单独创建池并将其直接传递给 map()。

  注意

  在提供低值种子时应谨慎，因为这可能会导致实验之间的自相关，例如：如果请求 8 个工作进程且种子为 4，则工作进程种子范围为 4 到 11。如果种子为 5，则工作进程种子范围为 5 到 12。 这两个实验将有 7 个种子的重叠，这可能会对结果产生意想不到的影响。

  注意

  为工作进程播种的目的是使每个工作进程都具有独立的种子，而不是在 map 方法的调用之间获得可重现的结果。换句话说，两个实验可能并且很可能会返回不同的结果，因为不可能知道哪个工作进程将选择哪个作业。 但是，我们可以确保每个工作进程都具有不同的种子，并且每个工作进程上的伪随机操作将是不相关的。

• max_tasks_per_child (int, 可选) – 每个子进程选取的最**大**作业数。默认为 None，即对作业数没有限制。

• worker_threads (int, 可选) – 工作进程的线程数。默认为 1。

• index_with_generator (bool, 可选) –

  如果为 True，则 tensordict 的拆分/分块将在查询期间完成，从而节省初始化时间。请注意，chunk() 和 split() 比索引（在生成器中使用的）效率更高，因此在初始化时节省处理时间可能会对总运行时产生负面影响。默认为 True。

  注意

  index_with_generator 的默认值对于 map_iter 和 map 不同，前者假定存储 TensorDict 的拆分版本在内存中成本过高。

• pbar (bool, 可选) – 如果为 True，将显示进度条。需要安装 tqdm。 默认为 False。

• mp_start_method (str, 可选) – 多进程的启动方法。如果未提供，将使用默认启动方法。接受的字符串为 "fork" 和 "spawn"。 请记住，"cuda" 张量不能在使用 "fork" 启动方法的进程之间共享。如果 pool 传递给 map 方法，则这不起作用。

示例

>>> import torch
>>> from tensordict import TensorDict
>>>
>>> def process_data(data):
...     data.unlock_()
...     data.set("y", data.get("x") + 1)
...     return data
>>> if __name__ == "__main__":
...     data = TensorDict({"x": torch.zeros(1, 1_000_000)}, [1, 1_000_000]).memmap_()
...     for sample in data.map_iter(process_data, dim=1, chunksize=5):
...         print(sample["y"])
...         break
...
tensor([[1., 1., 1., 1., 1.]])

masked_fill(mask: Tensor, value: float | bool) → T

masked_fill 的异位版本。

参数:

• mask (boolean torch.Tensor) – 要填充的值的掩码。形状必须与 tensordict 批次大小匹配。

• value – 用于填充张量的值。

返回:

self

示例

>>> td = TensorDict(source={'a': torch.zeros(3, 4)},
...     batch_size=[3])
>>> mask = torch.tensor([True, False, False])
>>> td1 = td.masked_fill(mask, 1.0)
>>> td1.get("a")
tensor([[1., 1., 1., 1.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]])

masked_fill_(mask: Tensor, value: float | bool) → T

使用所需值填充与掩码对应的值。

参数:

• mask (boolean torch.Tensor) – 要填充的值的掩码。形状必须与 tensordict 批次大小匹配。

• value – 用于填充张量的值。

返回:

self

示例

>>> td = TensorDict(source={'a': torch.zeros(3, 4)},
...     batch_size=[3])
>>> mask = torch.tensor([True, False, False])
>>> td.masked_fill_(mask, 1.0)
>>> td.get("a")
tensor([[1., 1., 1., 1.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]])

masked_select(mask: Tensor) → T

屏蔽 TensorDict 的所有张量，并返回一个新的 TensorDict 实例，其中类似的键指向屏蔽的值。

参数:

mask (torch.Tensor) – 用于张量的布尔掩码。形状必须与 TensorDict batch_size 匹配。

示例

>>> td = TensorDict(source={'a': torch.zeros(3, 4)},
...    batch_size=[3])
>>> mask = torch.tensor([True, False, False])
>>> td_mask = td.masked_select(mask)
>>> td_mask.get("a")
tensor([[0., 0., 0., 0.]])

maximum(other: tensordict.base.TensorDictBase | torch.Tensor, *, default: Optional[Union[str, Tensor]] = None) → T

计算 self 和 other 的元素级最大值。

参数:

other (TensorDict 或 Tensor) – 其他输入 tensordict 或张量。

关键字参数:

default (torch.Tensor 或 str, 可选) – 用于独占条目的默认值。如果未提供，则两个 tensordict 的键列表必须完全匹配。如果传递了 default="intersection"，则仅考虑相交的键集，其他键将被忽略。在所有其他情况下，default 将用于操作两侧的所有缺失条目。

maximum_(other: tensordict.base.TensorDictBase | torch.Tensor) → T

maximum() 的原地版本。

注意

原地 maximum 不支持 default 关键字参数。

classmethod maybe_dense_stack(items: Sequence[TensorDictBase], dim: int = 0, out: Optional[T] = None, strict: bool = False) → T

如果可能，密集堆叠张量或 tensordict，否则堆叠到 LazyStackedTensorDict 上。

示例

>>> td0 = TensorDict({"a": 0}, [])
>>> td1 = TensorDict({"b": 0}, [])
>>> LazyStackedTensorDict.maybe_dense_stack([td0, td0])  # returns a TensorDict with shape [2]
>>> LazyStackedTensorDict.maybe_dense_stack([td0, td1])  # returns a LazyStackedTensorDict with shape [2]
>>> LazyStackedTensorDict.maybe_dense_stack(list(torch.randn(2)))  # returns a torch.Tensor with shape [2]

mean(dim: Union[int, Tuple[int]] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, dtype: Optional[dtype] = None, reduce: Optional[bool] = None) → tensordict.base.TensorDictBase | torch.Tensor

返回输入 tensordict 中所有元素的平均值。

参数:

• dim (int, int 元组, 可选) – 如果为 None，则返回一个无量纲的 tensordict，其中包含所有叶子的平均值（如果可以计算）。如果为整数或整数元组，则仅当此维度与 tensordict 形状兼容时，才对指定的维度调用 mean。

• keepdim (bool) – 输出张量是否保留维度。

关键字参数:

• dtype (torch.dtype, 可选) – 返回张量的所需数据类型。如果指定，则在执行操作之前将输入张量转换为 dtype。这对于防止数据类型溢出很有用。默认值：None。

• reduce (bool, 可选) – 如果为 True，则将在所有 TensorDict 值上进行归约，并返回单个归约张量。默认为 False。

memmap(prefix: Optional[str] = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False, existsok: bool = True) → T

将所有张量写入新 tensordict 中相应的内存映射张量。

参数:

• prefix (str) – 内存映射张量将存储在其中的目录前缀。目录树结构将模仿 tensordict 的结构。

• copy_existing (bool) – 如果为 False（默认值），则当 tensordict 中的条目已经是存储在磁盘上且具有关联文件的张量，但未根据前缀保存在正确的位置时，将引发异常。如果为 True，则任何现有张量都将被复制到新位置。

关键字参数:

• num_threads (int, 可选) – 用于写入内存映射张量的线程数。默认为 0。

• return_early (bool, 可选) – 如果 True 且 num_threads>0，则该方法将返回 tensordict 的 future。

• share_non_tensor (bool, 可选) – 如果为 True，则非张量数据将在进程和写入操作之间共享（例如，在单个节点内的任何 worker 上进行原地更新或设置），并且将更新所有其他 worker 上的值。如果非张量叶子的数量很高（例如，共享大量非张量数据堆栈），则可能会导致 OOM 或类似错误。默认为 False。

• existsok (bool, 可选) – 如果为 False，则如果张量已存在于同一路径中，则会引发异常。默认为 True。

然后 TensorDict 将被锁定，这意味着任何非原地写入操作都将引发异常（例如，重命名、设置或删除条目）。一旦 tensordict 被解锁，内存映射属性将变为 False，因为不再保证跨进程标识。

返回:

如果 return_early=False，则返回一个新 tensordict，其中的张量存储在磁盘上；否则返回一个 TensorDictFuture 实例。

注意

以这种方式序列化对于深度嵌套的 tensordict 可能会很慢，因此不建议在训练循环内调用此方法。

memmap_(prefix: Optional[str] = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False, existsok: bool = True) → T

将所有张量写入相应的内存映射张量，原地操作。

参数:

• prefix (str) – 内存映射张量将存储在其中的目录前缀。目录树结构将模仿 tensordict 的结构。

• copy_existing (bool) – 如果为 False（默认值），则当 tensordict 中的条目已经是存储在磁盘上且具有关联文件的张量，但未根据前缀保存在正确的位置时，将引发异常。如果为 True，则任何现有张量都将被复制到新位置。

关键字参数:

• num_threads (int, 可选) – 用于写入内存映射张量的线程数。默认为 0。

• return_early (bool, optional) – if True and num_threads>0, the method will return a future of the tensordict. The resulting tensordict can be queried using future.result().

• share_non_tensor (bool, 可选) – 如果为 True，则非张量数据将在进程和写入操作之间共享（例如，在单个节点内的任何 worker 上进行原地更新或设置），并且将更新所有其他 worker 上的值。如果非张量叶子的数量很高（例如，共享大量非张量数据堆栈），则可能会导致 OOM 或类似错误。默认为 False。

• existsok (bool, 可选) – 如果为 False，则如果张量已存在于同一路径中，则会引发异常。默认为 True。

然后 TensorDict 将被锁定，这意味着任何非原地写入操作都将引发异常（例如，重命名、设置或删除条目）。一旦 tensordict 被解锁，内存映射属性将变为 False，因为不再保证跨进程标识。

返回:

如果 return_early=False，则返回 self；否则返回 TensorDictFuture 实例。

注意

以这种方式序列化对于深度嵌套的 tensordict 可能会很慢，因此不建议在训练循环内调用此方法。

memmap_like(prefix: Optional[str] = None, copy_existing: bool = False, *, existsok: bool = True, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False) → T

创建一个无内容的内存映射 tensordict，其形状与原始 tensordict 相同。

参数:

• prefix (str) – 内存映射张量将存储在其中的目录前缀。目录树结构将模仿 tensordict 的结构。

• copy_existing (bool) – 如果为 False（默认值），则当 tensordict 中的条目已经是存储在磁盘上且具有关联文件的张量，但未根据前缀保存在正确的位置时，将引发异常。如果为 True，则任何现有张量都将被复制到新位置。

关键字参数:

• num_threads (int, 可选) – 用于写入内存映射张量的线程数。默认为 0。

• return_early (bool, 可选) – 如果 True 且 num_threads>0，则该方法将返回 tensordict 的 future。

• share_non_tensor (bool, 可选) – 如果为 True，则非张量数据将在进程和写入操作之间共享（例如，在单个节点内的任何 worker 上进行原地更新或设置），并且将更新所有其他 worker 上的值。如果非张量叶子的数量很高（例如，共享大量非张量数据堆栈），则可能会导致 OOM 或类似错误。默认为 False。

• existsok (bool, 可选) – 如果为 False，则如果张量已存在于同一路径中，则会引发异常。默认为 True。

然后 TensorDict 将被锁定，这意味着任何非原地写入操作都将引发异常（例如，重命名、设置或删除条目）。一旦 tensordict 被解锁，内存映射属性将变为 False，因为不再保证跨进程标识。

返回:

如果 return_early=False，则返回一个新的 TensorDict 实例，其中数据存储为内存映射张量；否则返回一个 TensorDictFuture 实例。

注意

这是在磁盘上写入一组大型缓冲区的推荐方法，因为 memmap_() 将复制信息，对于大型内容来说，这可能会很慢。

示例

>>> td = TensorDict({
...     "a": torch.zeros((3, 64, 64), dtype=torch.uint8),
...     "b": torch.zeros(1, dtype=torch.int64),
... }, batch_size=[]).expand(1_000_000)  # expand does not allocate new memory
>>> buffer = td.memmap_like("/path/to/dataset")

memmap_refresh_()

如果内存映射 tensordict 具有 saved_path，则刷新其内容。

如果未与其关联路径，则此方法将引发异常。

minimum(other: tensordict.base.TensorDictBase | torch.Tensor, *, default: Optional[Union[str, Tensor]] = None) → T

计算 self 和 other 的元素级最小值。

参数:

other (TensorDict 或 Tensor) – 其他输入 tensordict 或张量。

关键字参数:

default (torch.Tensor 或 str, 可选) – 用于独占条目的默认值。如果未提供，则两个 tensordict 的键列表必须完全匹配。如果传递了 default="intersection"，则仅考虑相交的键集，其他键将被忽略。在所有其他情况下，default 将用于操作两侧的所有缺失条目。

minimum_(other: tensordict.base.TensorDictBase | torch.Tensor) → T

minimum() 的原地版本。

注意

原地 minimum 不支持 default 关键字参数。

mul(other: tensordict.base.TensorDictBase | torch.Tensor, *, default: Optional[Union[str, Tensor]] = None) → T

将 other 乘以 self。

\[\text{{out}}_i = \text{{input}}_i \times \text{{other}}_i\]

支持广播、类型提升以及整数、浮点数和复数输入。

参数:

other (TensorDict, Tensor 或 Number) – 要从 self 中减去的张量或数字。

关键字参数:

default (torch.Tensor 或 str, 可选) – 用于独占条目的默认值。如果未提供，则两个 tensordict 的键列表必须完全匹配。如果传递了 default="intersection"，则仅考虑相交的键集，其他键将被忽略。在所有其他情况下，default 将用于操作两侧的所有缺失条目。

mul_(other: tensordict.base.TensorDictBase | torch.Tensor) → T

mul() 的原地版本。

注意

原地 mul 不支持 default 关键字参数。

named_apply(fn: Callable, *others: T, nested_keys: bool = False, batch_size: Optional[Sequence[int]] = None, device: torch.device | None = _NoDefault.ZERO, names: Optional[Sequence[str]] = _NoDefault.ZERO, inplace: bool = False, default: Any = _NoDefault.ZERO, filter_empty: Optional[bool] = None, propagate_lock: bool = False, call_on_nested: bool = False, out: Optional[TensorDictBase] = None, **constructor_kwargs) → Optional[T]

将键条件可调用对象应用于 tensordict 中存储的所有值，并将它们设置在一个新的 atensordict 中。

可调用签名必须是 Callable[Tuple[str, Tensor, ...], Optional[Union[Tensor, TensorDictBase]]]。

参数:

• fn (Callable) – 要应用于 tensordict 中的 (name, tensor) 对的函数。对于每个叶节点，将仅使用其叶节点名称（而不是完整的 NestedKey）。

• *others (TensorDictBase 实例, 可选) – 如果提供，这些 tensordict 实例应具有与 self 结构匹配的结构。fn 参数应接收与 tensordict 数量（包括 self）一样多的未命名输入。如果其他 tensordict 缺少条目，则可以通过 default 关键字参数传递默认值。

• nested_keys (bool, optional) – 如果为 True，将使用到叶节点的完整路径。默认为 False，即仅将最后一个字符串传递给函数。

• batch_size (int 序列, 可选) – 如果提供，结果 TensorDict 将具有所需的 batch_size。batch_size 参数应与转换后的 batch_size 匹配。这是一个仅关键字参数。

• device (torch.device, 可选) – 结果设备，如果有。

• names (str 列表, 可选) – 新的维度名称，以防 batch_size 被修改。

• inplace (bool, optional) – 如果为 True，则进行原地更改。默认为 False。这是一个仅关键字参数。

• default (Any, 可选) – 其他 tensordict 中缺少条目的默认值。如果未提供，缺少条目将引发 KeyError。

• filter_empty (bool, optional) – 如果为 True，将过滤掉空 tensordict。这也降低了计算成本，因为不会创建和销毁空数据结构。为了向后兼容，默认为 False。

• propagate_lock (bool, optional) – 如果为 True，锁定的 tensordict 将生成另一个锁定的 tensordict。默认为 False。

• call_on_nested (bool, optional) –

  如果 True，则将对第一级张量和容器（TensorDict 或 tensorclass）调用该函数。在这种情况下，func 负责将其调用传播到嵌套级别。这允许在将调用传播到嵌套 tensordict 时进行细粒度行为。如果 False，则仅在叶节点上调用该函数，并且 apply 将负责将函数分派到所有叶节点。

  >>> td = TensorDict({"a": {"b": [0.0, 1.0]}, "c": [1.0, 2.0]})
  >>> def mean_tensor_only(val):
  ...     if is_tensor_collection(val):
  ...         raise RuntimeError("Unexpected!")
  ...     return val.mean()
  >>> td_mean = td.apply(mean_tensor_only)
  >>> def mean_any(val):
  ...     if is_tensor_collection(val):
  ...         # Recurse
  ...         return val.apply(mean_any, call_on_nested=True)
  ...     return val.mean()
  >>> td_mean = td.apply(mean_any, call_on_nested=True)
  

• out (TensorDictBase, 可选) –

  用于写入结果的 tensordict。这可以用于避免创建新的 tensordict

  >>> td = TensorDict({"a": 0})
  >>> td.apply(lambda x: x+1, out=td)
  >>> assert (td==1).all()
  

  警告

  如果对 tensordict 执行的操作需要访问多个键才能进行单个计算，则提供等于 self 的 out 参数可能会导致操作提供静默的错误结果。例如

  >>> td = TensorDict({"a": 1, "b": 1})
  >>> td.apply(lambda x: x+td["a"])["b"] # Right!
  tensor(2)
  >>> td.apply(lambda x: x+td["a"], out=td)["b"] # Wrong!
  tensor(3)
  

• **constructor_kwargs – 要传递给 TensorDict 构造函数的其他关键字参数。

返回:

具有 transformed_in 张量的新 tensordict。

示例

>>> td = TensorDict({
...     "a": -torch.ones(3),
...     "nested": {"a": torch.ones(3), "b": torch.zeros(3)}},
...     batch_size=[3])
>>> def name_filter(name, tensor):
...     if name == "a":
...         return tensor
>>> td.named_apply(name_filter)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
        nested: TensorDict(
            fields={
                a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3]),
    device=None,
    is_shared=False)
>>> def name_filter(name, *tensors):
...     if name == "a":
...         r = 0
...         for tensor in tensors:
...             r = r + tensor
...         return tensor
>>> out = td.named_apply(name_filter, td)
>>> print(out)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
        nested: TensorDict(
            fields={
                a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3]),
    device=None,
    is_shared=False)
>>> print(out["a"])
tensor([-1., -1., -1.])

注意

如果函数返回 None，则忽略该条目。这可以用于过滤 tensordict 中的数据

>>> td = TensorDict({"1": 1, "2": 2, "b": {"2": 2, "1": 1}}, [])
>>> def name_filter(name, tensor):
...     if name == "1":
...         return tensor
>>> td.named_apply(name_filter)
TensorDict(
    fields={
        1: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        b: TensorDict(
            fields={
                1: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

property names

tensordict 的维度名称。

名称可以在构造时使用 names 参数设置。

另请参阅 refine_names() 以了解有关如何在构造后设置名称的详细信息。

nanmean(dim: Union[int, Tuple[int]] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, dtype: Optional[dtype] = None, reduce: Optional[bool] = None) → tensordict.base.TensorDictBase | torch.Tensor

返回输入 tensordict 中所有非 NaN 元素的平均值。

参数:

• dim (int, int 元组, 可选) – 如果为 None，则返回一个无量纲的 tensordict，其中包含所有叶子的平均值（如果可以计算）。如果为整数或整数元组，则仅当此维度与 tensordict 形状兼容时，才对指定的维度调用 mean。

• keepdim (bool) – 输出张量是否保留维度。

关键字参数:

• dtype (torch.dtype, 可选) – 返回张量的所需数据类型。如果指定，则在执行操作之前将输入张量转换为 dtype。这对于防止数据类型溢出很有用。默认值：None。

• reduce (bool, 可选) – 如果为 True，则将在所有 TensorDict 值上进行归约，并返回单个归约张量。默认为 False。

nansum(dim: Union[int, Tuple[int]] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, dtype: Optional[dtype] = None, reduce: Optional[bool] = None) → tensordict.base.TensorDictBase | torch.Tensor

返回输入 tensordict 中所有非 NaN 元素的总和。

参数:

• dim (int, int 元组, 可选) – 如果为 None，则返回一个无量纲的 tensordict，其中包含所有叶节点的总和值（如果可以计算）。如果为整数或整数元组，则仅当此维度与 tensordict 形状兼容时，才对指定的维度调用 sum。

• keepdim (bool) – 输出张量是否保留维度。

关键字参数:

• dtype (torch.dtype, 可选) – 返回张量的所需数据类型。如果指定，则在执行操作之前将输入张量转换为 dtype。这对于防止数据类型溢出很有用。默认值：None。

• reduce (bool, 可选) – 如果为 True，则将在所有 TensorDict 值上进行归约，并返回单个归约张量。默认为 False。

property ndim: int

See batch_dims().

ndimension() → int

See batch_dims().

neg() → T

计算 TensorDict 中每个元素的 neg() 值。

neg_() → T

原地计算 TensorDict 中每个元素的 neg() 值。

new_empty(*size: Size, dtype: Optional[dtype] = None, device: device, str, int] = _NoDefault.ZERO, requires_grad: bool = False, layout: layout = torch.strided, pin_memory: Optional[bool] = None)

返回大小为 size 且包含空张量的 TensorDict。

默认情况下，返回的 TensorDict 具有与此 tensordict 相同的 torch.dtype 和 torch.device。

参数:

size (int...) – 定义输出张量形状的整数列表、元组或 torch.Size。

关键字参数:

• dtype (torch.dtype, optional) – 返回的 tensordict 的所需类型。默认值：如果为 None，则 torch.dtype 将保持不变。

• device (torch.device, optional) – 返回的 tensordict 的所需设备。默认值：如果为 None，则 torch.device 将保持不变。

• requires_grad (bool, optional) – 如果 autograd 应该记录对返回张量的操作。默认值：False。

• layout (torch.layout, optional) – 返回的 TensorDict 值的所需布局。默认值：torch.strided。

• pin_memory (bool, optional) – 如果设置，则返回的张量将分配在锁页内存中。仅适用于 CPU 张量。默认值：False。

new_full(size: Size, fill_value, *, dtype: Optional[dtype] = None, device: device, str, int] = _NoDefault.ZERO, requires_grad: bool = False, layout: layout = torch.strided, pin_memory: Optional[bool] = None)

返回大小为 size 并填充了 1 的 TensorDict。

默认情况下，返回的 TensorDict 具有与此 tensordict 相同的 torch.dtype 和 torch.device。

参数:

• size (整数序列) – 定义输出张量形状的整数列表、元组或 torch.Size。

• fill_value (标量) – 用于填充输出张量的数字。

关键字参数:

• dtype (torch.dtype, optional) – 返回的 tensordict 的所需类型。默认值：如果为 None，则 torch.dtype 将保持不变。

• device (torch.device, optional) – 返回的 tensordict 的所需设备。默认值：如果为 None，则 torch.device 将保持不变。

• requires_grad (bool, optional) – 如果 autograd 应该记录对返回张量的操作。默认值：False。

• layout (torch.layout, optional) – 返回的 TensorDict 值的所需布局。默认值：torch.strided。

• pin_memory (bool, optional) – 如果设置，则返回的张量将分配在锁页内存中。仅适用于 CPU 张量。默认值：False。

new_ones(*size: Size, dtype: Optional[dtype] = None, device: Union[device, str, int] = _NoDefault.ZERO, requires_grad: bool = False, layout: layout = torch.strided, pin_memory: Optional[bool] = None)

返回大小为 size 并填充了 1 的 TensorDict。

默认情况下，返回的 TensorDict 具有与此 tensordict 相同的 torch.dtype 和 torch.device。

参数:

size (int...) – 定义输出张量形状的整数列表、元组或 torch.Size。

关键字参数:

• dtype (torch.dtype, optional) – 返回的 tensordict 的所需类型。默认值：如果为 None，则 torch.dtype 将保持不变。

• device (torch.device, optional) – 返回的 tensordict 的所需设备。默认值：如果为 None，则 torch.device 将保持不变。

• requires_grad (bool, optional) – 如果 autograd 应该记录对返回张量的操作。默认值：False。

• layout (torch.layout, optional) – 返回的 TensorDict 值的所需布局。默认值：torch.strided。

• pin_memory (bool, optional) – 如果设置，则返回的张量将分配在锁页内存中。仅适用于 CPU 张量。默认值：False。

new_tensor(data: torch.Tensor | tensordict.base.TensorDictBase, *, dtype: Optional[dtype] = None, device: Union[device, str, int] = _NoDefault.ZERO, requires_grad: bool = False, pin_memory: Optional[bool] = None)

返回一个新的 TensorDict，其数据为张量 data。

默认情况下，返回的 TensorDict 值与此张量具有相同的 torch.dtype 和 torch.device。

data 也可以是张量集合（TensorDict 或 tensorclass），在这种情况下，new_tensor 方法会遍历 self 和 data 的张量对。

参数:

data (torch.Tensor 或 TensorDictBase) – 要复制的数据。

关键字参数:

• dtype (torch.dtype, optional) – 返回的 tensordict 的所需类型。默认值：如果为 None，则 torch.dtype 将保持不变。

• device (torch.device, optional) – 返回的 tensordict 的所需设备。默认值：如果为 None，则 torch.device 将保持不变。

• requires_grad (bool, optional) – 如果 autograd 应该记录对返回张量的操作。默认值：False。

• pin_memory (bool, optional) – 如果设置，则返回的张量将分配在锁页内存中。仅适用于 CPU 张量。默认值：False。

new_zeros(*size: Size, dtype: Optional[dtype] = None, device: Union[device, str, int] = _NoDefault.ZERO, requires_grad: bool = False, layout: layout = torch.strided, pin_memory: Optional[bool] = None)

返回一个大小为 size 并填充 0 的 TensorDict。

默认情况下，返回的 TensorDict 具有与此 tensordict 相同的 torch.dtype 和 torch.device。

参数:

size (int...) – 定义输出张量形状的整数列表、元组或 torch.Size。

关键字参数:

• dtype (torch.dtype, optional) – 返回的 tensordict 的所需类型。默认值：如果为 None，则 torch.dtype 将保持不变。

• device (torch.device, optional) – 返回的 tensordict 的所需设备。默认值：如果为 None，则 torch.device 将保持不变。

• requires_grad (bool, optional) – 如果 autograd 应该记录对返回张量的操作。默认值：False。

• layout (torch.layout, optional) – 返回的 TensorDict 值的所需布局。默认值：torch.strided。

• pin_memory (bool, optional) – 如果设置，则返回的张量将分配在锁页内存中。仅适用于 CPU 张量。默认值：False。

non_tensor_items(include_nested: bool = False)

返回所有非张量叶节点，可能是递归地返回。

norm(*, out=None, dtype: torch.dtype | None = None)

计算 tensordict 中每个张量的范数。

关键字参数:

• out (TensorDict, 可选) – 输出 tensordict。

• dtype (torch.dtype, 可选) – 输出 dtype (torch>=2.4)。

numel() → int

批次中的元素总数。

下限为 1，因为两个空形状 tensordict 的堆叠将有两个元素，因此我们认为 tensordict 至少是 1 个元素大小。

numpy()

将 tensordict 转换为 numpy 数组的（可能是嵌套的）字典。

非张量数据按原样公开。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> data = TensorDict({"a": {"b": torch.zeros(()), "c": "a string!"}})
>>> print(data)
TensorDict(
    fields={
        a: TensorDict(
            fields={
                b: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                c: NonTensorData(data=a string!, batch_size=torch.Size([]), device=None)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> print(data.numpy())
{'a': {'b': array(0., dtype=float32), 'c': 'a string!'}}

param_count(*, count_duplicates: bool = True) → int

计算参数的数量（可索引项的总数），仅计算张量。

关键字参数:

count_duplicates (bool) – 是否将重复张量计为独立张量。如果为 False，则仅会丢弃严格相同的张量（来自公共基础张量的相同视图但具有不同 id 的张量将被计数两次）。默认为 True（每个张量都假定为单个副本）。

permute(*args, **kwargs)

返回一个 tensordict 的视图，其批次维度根据 dims 进行置换。

参数:

• *dims_list (int) – tensordict 的批次维度的新顺序。 或者，可以提供单个整数迭代器。

• dims (int 列表) – 调用 permute(…) 的替代方法。

返回:

一个新的 tensordict，其批次维度按所需顺序排列。

示例

>>> tensordict = TensorDict({"a": torch.randn(3, 4, 5)}, [3, 4])
>>> print(tensordict.permute([1, 0]))
PermutedTensorDict(
    source=TensorDict(
        fields={
            a: Tensor(torch.Size([3, 4, 5]), dtype=torch.float32)},
        batch_size=torch.Size([3, 4]),
        device=cpu,
        is_shared=False),
    op=permute(dims=[1, 0]))
>>> print(tensordict.permute(1, 0))
PermutedTensorDict(
    source=TensorDict(
        fields={
            a: Tensor(torch.Size([3, 4, 5]), dtype=torch.float32)},
        batch_size=torch.Size([3, 4]),
        device=cpu,
        is_shared=False),
    op=permute(dims=[1, 0]))
>>> print(tensordict.permute(dims=[1, 0]))
PermutedTensorDict(
    source=TensorDict(
        fields={
            a: Tensor(torch.Size([3, 4, 5]), dtype=torch.float32)},
        batch_size=torch.Size([3, 4]),
        device=cpu,
        is_shared=False),
    op=permute(dims=[1, 0]))

pin_memory(num_threads: Optional[int] = None, inplace: bool = False) → T

在存储的张量上调用 pin_memory()。

参数:

• num_threads (int 或 str) – 如果提供，则用于在叶节点上调用 pin_memory 的线程数。 默认为 None，这会在 ThreadPoolExecutor(max_workers=None) 中设置大量线程。 要在主线程上执行所有对 pin_memory() 的调用，请传递 num_threads=0。

• inplace (bool, 可选) – 如果为 True，则 tensordict 将就地修改。 默认为 False。

pin_memory_(num_threads: int | str = 0) → T

在存储的张量上调用 pin_memory()，并返回就地修改的 TensorDict。

参数:

num_threads (int 或 str) – 如果提供，则用于在叶节点上调用 pin_memory 的线程数。 如果传递 "auto"，则会自动确定线程数。

pop(key: NestedKey, default: Any = _NoDefault.ZERO) → Tensor

从 tensordict 中移除并返回一个值。

如果该值不存在且未提供默认值，则会引发 KeyError。

参数:

• key (str 或 nested key) – 要查找的条目。

• default (Any, 可选) – 如果找不到键，则返回的值。

示例

>>> td = TensorDict({"1": 1}, [])
>>> one = td.pop("1")
>>> assert one == 1
>>> none = td.pop("1", default=None)
>>> assert none is None

popitem() → Tuple[NestedKey, Tensor]

移除最后插入 TensorDict 的项目。

popitem 将仅返回非嵌套值。

pow(other: tensordict.base.TensorDictBase | torch.Tensor, *, default: Optional[Union[str, Tensor]] = None) → T

计算 self 中每个元素与 other 的幂，并返回包含结果的张量。

other 可以是单个 float 数字、Tensor 或 TensorDict。

当 other 是张量时，input 和 other 的形状必须是可广播的。

参数:

other (float, tensor 或 tensordict) – 指数值

关键字参数:

default (torch.Tensor 或 str, 可选) – 用于独占条目的默认值。如果未提供，则两个 tensordict 的键列表必须完全匹配。如果传递了 default="intersection"，则仅考虑相交的键集，其他键将被忽略。在所有其他情况下，default 将用于操作两侧的所有缺失条目。

pow_(other: tensordict.base.TensorDictBase | torch.Tensor) → T

pow() 的原地版本。

注意

原地 pow 不支持 default 关键字参数。

prod(dim: Union[int, Tuple[int]] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, dtype: Optional[dtype] = None, reduce: Optional[bool] = None) → tensordict.base.TensorDictBase | torch.Tensor

返回输入 tensordict 中所有元素的乘积值。

参数:

• dim (int, int 元组, 可选) – 如果为 None，则返回一个无量纲的 tensordict，其中包含所有叶节点的 prod 值（如果可以计算）。 如果为整数或整数元组，则仅当此维度与 tensordict 形状兼容时，才对指定的维度调用 prod。

• keepdim (bool) – 输出张量是否保留维度。

关键字参数:

• dtype (torch.dtype, 可选) – 返回张量的所需数据类型。如果指定，则在执行操作之前将输入张量转换为 dtype。这对于防止数据类型溢出很有用。默认值：None。

• reduce (bool, 可选) – 如果为 True，则将在所有 TensorDict 值上进行归约，并返回单个归约张量。默认为 False。

qint32()

将所有张量转换为 torch.qint32。

qint8()

将所有张量转换为 torch.qint8。

quint4x2()

将所有张量转换为 torch.quint4x2。

quint8()

将所有张量转换为 torch.quint8。

reciprocal() → T

计算 TensorDict 中每个元素的 reciprocal() 值。

reciprocal_() → T

就地计算 TensorDict 中每个元素的 reciprocal() 值。

record_stream(stream: Stream)

将 tensordict 标记为已被此流使用。

当 tensordict 被释放时，确保在释放时流上排队的所有工作完成之前，张量内存不会被其他张量重用。

有关更多信息，请参阅 record_stream()。

recv(src: int, *, group: 'dist.ProcessGroup' | None = None, init_tag: int = 0, pseudo_rand: bool = False) → int

接收 tensordict 的内容并使用其内容更新当前 tensordict。

请查看 send 方法中的示例以了解上下文。

参数:

src (int) – 源工作进程的 rank。

关键字参数:

• group (torch.distributed.ProcessGroup, 可选) – 如果设置，则指定的进程组将用于通信。否则，将使用默认进程组。默认为 None。

• init_tag (int) – 源工作进程使用的 init_tag。

• pseudo_rand (bool) – 如果为 True，标签序列将是伪随机的，允许从不同节点发送多个数据而不会重叠。 请注意，生成这些伪随机数是昂贵的（1e-5 秒/数字），这意味着它可能会减慢算法的运行时速度。 此值必须与传递给 send() 的值匹配。 默认为 False。

reduce(dst, op=None, async_op=False, return_premature=False, group=None)

跨所有机器规约 tensordict。

只有 rank 为 dst 的进程会收到最终结果。

refine_names(*names) → T

根据名称优化自身的维度名称。

优化是重命名的一种特殊情况，它可以“提升”未命名的维度。可以将 None 维度优化为具有任何名称；命名的维度只能优化为具有相同的名称。

由于命名张量可以与未命名张量共存，因此优化名称提供了一种编写命名张量感知代码的良好方法，该代码可用于命名张量和未命名张量。

names 最多可以包含一个省略号 (...)。省略号会贪婪地展开；它会就地展开，以使用 self.names 中相应索引的名称填充 names，使其长度与 self.dim() 相同。

返回: 根据输入命名维度的相同 tensordict。

示例

>>> td = TensorDict({}, batch_size=[3, 4, 5, 6])
>>> tdr = td.refine_names(None, None, None, "d")
>>> assert tdr.names == [None, None, None, "d"]
>>> tdr = td.refine_names("a", None, None, "d")
>>> assert tdr.names == ["a", None, None, "d"]

rename(*names, **rename_map)

返回维度已重命名的 tensordict 的克隆。

示例

>>> td = TensorDict({}, batch_size=[1, 2, 3 ,4])
>>> td.names = list("abcd")
>>> td_rename = td.rename(c="g")
>>> assert td_rename.names == list("abgd")

rename_(*names, **rename_map)

与 rename() 相同，但就地执行重命名。

示例

>>> td = TensorDict({}, batch_size=[1, 2, 3 ,4])
>>> td.names = list("abcd")
>>> assert td.rename_(c="g")
>>> assert td.names == list("abgd")

rename_key_(old_key: NestedKey, new_key: NestedKey, safe: bool = False) → T

使用新字符串重命名键，并返回具有更新键名称的相同 tensordict。

参数:

• old_key (str 或 嵌套键) – 要重命名的键。

• new_key (str 或 嵌套键) – 条目的新名称。

• safe (bool, 可选) – 如果 True，当新键已存在于 TensorDict 中时，将抛出错误。

返回:

self

replace(*args, **kwargs)

创建 tensordict 的浅拷贝，其中的条目已被替换。

接受一个未命名的参数，该参数必须是 TensorDictBase 子类的字典。 此外，可以使用命名关键字参数更新第一级条目。

返回:

如果输入非空，则返回 self 的副本，其中条目已更新。 如果提供空字典或未提供字典，并且 kwargs 为空，则返回 self。

requires_grad_(requires_grad=True) → T

更改 autograd 是否应记录此张量上的操作：就地设置此张量的 requires_grad 属性。

返回此 tensordict。

参数:

requires_grad (bool, 可选) – autograd 是否应记录此 tensordict 上的操作。 默认为 True。

reshape(*args, **kwargs) → T

返回所需形状的连续的、重塑的张量。

参数:

*shape (int) – 结果 tensordict 的新形状。

返回:

具有重塑键的 TensorDict

示例

>>> td = TensorDict({
...     'x': torch.arange(12).reshape(3, 4),
... }, batch_size=[3, 4])
>>> td = td.reshape(12)
>>> print(td['x'])
torch.Tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11])

round() → T

计算 TensorDict 中每个元素的 round() 值。

round_() → T

就地计算 TensorDict 中每个元素的 round() 值。

save(prefix: Optional[str] = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False) → T

Saves the tensordict to disk.

This function is a proxy to memmap().

property saved_path

返回 memmap 保存的 TensorDict 的存储路径。

一旦 is_memmap() 返回 False（例如，当 tensordict 被解锁时），此参数将消失。

select(*keys: NestedKey, inplace: bool = False, strict: bool = True) → T

选择 tensordict 的键，并返回仅包含所选键的新 tensordict。

值不会被复制：对原始或新 tensordict 的张量进行就地修改将导致两个 tensordict 都发生更改。

参数:

• *keys (str) – 要选择的键

• inplace (bool) – 如果为 True，则就地修剪 tensordict。默认为 False。

• strict (bool, 可选) – 选择不存在的键是否会返回错误。 默认值： True。

返回:

新的 tensordict（如果 inplace=True，则为相同的 tensordict），仅包含所选键。

注意

要在 tensordict 中选择键并返回一个不包含这些键的 tensordict 版本，请参阅 split_keys() 方法。

示例

>>> from tensordict import TensorDict
>>> td = TensorDict({"a": 0, "b": {"c": 1, "d": 2}}, [])
>>> td.select("a", ("b", "c"))
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> td.select("a", "b")
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
                d: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> td.select("this key does not exist", strict=False)
TensorDict(
    fields={
    },
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

send(dst: int, *, group: 'dist.ProcessGroup' | None = None, init_tag: int = 0, pseudo_rand: bool = False) → None

将 tensordict 的内容发送到远程工作进程。

参数:

dst (int) – 目标工作进程的 rank，内容应发送到该进程。

关键字参数:

• group (torch.distributed.ProcessGroup, 可选) – 如果设置，则指定的进程组将用于通信。否则，将使用默认进程组。默认为 None。

• init_tag (int) – 用于标记张量的初始标签。请注意，这将递增 tensordict 中包含的张量数量。

• pseudo_rand (bool) – 如果为 True，则标签序列将是伪随机的，允许从不同节点发送多个数据而不会重叠。请注意，这些伪随机数的生成成本很高（1e-5 秒/数字），这意味着它可能会减慢算法的运行时速度。默认为 False。

示例

>>> from torch import multiprocessing as mp
>>> from tensordict import TensorDict
>>> import torch
>>>
>>>
>>> def client():
...     torch.distributed.init_process_group(
...         "gloo",
...         rank=1,
...         world_size=2,
...         init_method=f"tcp://127.0.0.1:10003",
...     )
...
...     td = TensorDict(
...         {
...             ("a", "b"): torch.randn(2),
...             "c": torch.randn(2, 3),
...             "_": torch.ones(2, 1, 5),
...         },
...         [2],
...     )
...     td.send(0)
...
>>>
>>> def server(queue):
...     torch.distributed.init_process_group(
...         "gloo",
...         rank=0,
...         world_size=2,
...         init_method=f"tcp://127.0.0.1:10003",
...     )
...     td = TensorDict(
...         {
...             ("a", "b"): torch.zeros(2),
...             "c": torch.zeros(2, 3),
...             "_": torch.zeros(2, 1, 5),
...         },
...         [2],
...     )
...     td.recv(1)
...     assert (td != 0).all()
...     queue.put("yuppie")
...
>>>
>>> if __name__=="__main__":
...     queue = mp.Queue(1)
...     main_worker = mp.Process(target=server, args=(queue,))
...     secondary_worker = mp.Process(target=client)
...
...     main_worker.start()
...     secondary_worker.start()
...     out = queue.get(timeout=10)
...     assert out == "yuppie"
...     main_worker.join()
...     secondary_worker.join()

set(key: NestedKey, item: Tensor, inplace: bool = False, *, non_blocking: bool = False, **kwargs: Any) → T

设置新的键值对。

参数:

• key (str, str 元组) – 要设置的键的名称。

• item (torch.Tensor 或 等效类型, TensorDictBase 实例) – 要存储在 tensordict 中的值。

• inplace (bool, 可选) – 如果 True 并且键与 tensordict 中的现有键匹配，则将对该键值对执行就地更新。 如果 inplace 为 True 且找不到条目，则会添加该条目。 对于更严格的就地操作，请改用 set_()。 默认为 False。

关键字参数:

non_blocking (bool, 可选) – 如果 True 且此复制发生在不同设备之间，则复制可能会相对于主机异步发生。

返回:

self

示例

>>> td = TensorDict({}, batch_size[3, 4])
>>> td.set("x", torch.randn(3, 4))
>>> y = torch.randn(3, 4, 5)
>>> td.set("y", y, inplace=True) # works, even if 'y' is not present yet
>>> td.set("y", torch.zeros_like(y), inplace=True)
>>> assert (y==0).all() # y values are overwritten
>>> td.set("y", torch.ones(5), inplace=True) # raises an exception as shapes mismatch

set_(key: NestedKey, item: Tensor, *, non_blocking: bool = False) → T

将值设置为现有键，同时保留原始存储。

参数:

• key (str) – 值的名称

• item (torch.Tensor 或 兼容类型, TensorDictBase) – 要存储在 tensordict 中的值

关键字参数:

non_blocking (bool, 可选) – 如果 True 且此复制发生在不同设备之间，则复制可能会相对于主机异步发生。

返回:

self

示例

>>> td = TensorDict({}, batch_size[3, 4])
>>> x = torch.randn(3, 4)
>>> td.set("x", x)
>>> td.set_("x", torch.zeros_like(x))
>>> assert (x == 0).all()

set_at_(key: NestedKey, value: Tensor, index: Union[None, int, slice, str, Tensor, List[Any], Tuple[Any, ...]], *, non_blocking: bool = False) → T

在 index 指示的索引处就地设置值。

参数:

• key (str, str 元组) – 要修改的键。

• value (torch.Tensor) – 要在索引 index 处设置的值

• index (int, 张量 或 元组) – 要写入值的索引。

关键字参数:

non_blocking (bool, 可选) – 如果 True 且此复制发生在不同设备之间，则复制可能会相对于主机异步发生。

返回:

self

示例

>>> td = TensorDict({}, batch_size[3, 4])
>>> x = torch.randn(3, 4)
>>> td.set("x", x)
>>> td.set_at_("x", value=torch.ones(1, 4), index=slice(1))
>>> assert (x[0] == 1).all()

set_non_tensor(key: NestedKey, value: Any)

使用 tensordict.tensorclass.NonTensorData 在 tensordict 中注册非张量值。

可以使用 TensorDictBase.get_non_tensor() 或直接使用 get 检索该值，这将返回 tensordict.tensorclass.NonTensorData 对象。

返回: self

示例

>>> data = TensorDict({}, batch_size=[])
>>> data.set_non_tensor(("nested", "the string"), "a string!")
>>> assert data.get_non_tensor(("nested", "the string")) == "a string!"
>>> # regular `get` works but returns a NonTensorData object
>>> data.get(("nested", "the string"))
NonTensorData(
    data='a string!',
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

setdefault(key: NestedKey, default: Tensor, inplace: bool = False) → Tensor

如果 tensordict 中没有 key，则插入 key 条目，其值为 default。

如果 tensordict 中有 key，则返回 key 的值，否则返回 default。

参数:

• key (str 或 嵌套键) – 值的名称。

• default (torch.Tensor 或 兼容类型, TensorDictBase) – 如果键尚不存在，则要存储在 tensordict 中的值。

返回:

tensordict 中键的值。 如果键之前未设置，则将为 default。

示例

>>> td = TensorDict({}, batch_size=[3, 4])
>>> val = td.setdefault("a", torch.zeros(3, 4))
>>> assert (val == 0).all()
>>> val = td.setdefault("a", torch.ones(3, 4))
>>> assert (val == 0).all() # output is still 0

property shape: Size

请参阅 batch_size。

share_memory_() → T

将所有张量放置在共享内存中。

TensorDict 随后会被锁定，这意味着任何非原地写入操作都会抛出异常（例如，重命名、设置或移除条目）。相反地，一旦 tensordict 被解锁，share_memory 属性会变为 False，因为跨进程的标识不再得到保证。

返回:

self

sigmoid() → T

计算 TensorDict 中每个元素的 sigmoid() 值。

sigmoid_() → T

原地计算 TensorDict 中每个元素的 sigmoid() 值。

sign() → T

计算 TensorDict 中每个元素的 sign() 值。

sign_() → T

原地计算 TensorDict 中每个元素的 sign() 值。

sin() → T

计算 TensorDict 中每个元素的 sin() 值。

sin_() → T

原地计算 TensorDict 中每个元素的 sin() 值。

sinh() → T

计算 TensorDict 中每个元素的 sinh() 值。

sinh_() → T

原地计算 TensorDict 中每个元素的 sinh() 值。

size(dim: Optional[int] = None) → torch.Size | int

返回由 dim 指示的维度的大小。

如果未指定 dim，则返回 TensorDict 的 batch_size 属性。

property sorted_keys: list[tensordict._nestedkey.NestedKey]

返回按字母顺序排序的键。

不支持额外的参数。

如果 TensorDict 被锁定，则键会被缓存，直到 tensordict 被解锁以实现更快的执行速度。

split(split_size: int | list[int], dim: int = 0) → list[tensordict.base.TensorDictBase]

在给定维度中，使用指定大小分割 TensorDict 中的每个张量，类似于 torch.split。

返回一个 TensorDict 实例列表，其中包含分割块的视图。

参数:

• split_size (int 或 List(int)) – 单个块的大小或每个块的大小列表。

• dim (int) – 沿其分割张量的维度。

返回:

一个 TensorDict 列表，其中包含给定维度中指定的大小。

示例

>>> td = TensorDict({
...     'x': torch.arange(12).reshape(3, 4),
... }, batch_size=[3, 4])
>>> td0, td1 = td.split([1, 2], dim=0)
>>> print(td0['x'])
torch.Tensor([[0, 1, 2, 3]])

split_keys(*key_sets, inplace=False, strict: bool = True, reproduce_struct: bool = False)

根据一个或多个键集将 tensordict 分割成子集。

该方法将返回 N+1 个 tensordict，其中 N 是提供的参数数量。

参数:

• inplace (bool, optional) – 如果为 True，则键会从 self 中原地移除。默认为 False。

• strict (bool, optional) – 如果为 True，当键丢失时会引发异常。默认为 True。

• reproduce_struct (bool, optional) – 如果为 True，则返回的所有 tensordict 都具有与 self 相同的树结构，即使某些子 tensordict 不包含叶节点。

注意

None 非张量值将被忽略且不返回。

注意

该方法不检查所提供列表中的重复项。

示例

>>> td = TensorDict(
...     a=0,
...     b=0,
...     c=0,
...     d=0,
... )
>>> td_a, td_bc, td_d = td.split_keys(["a"], ["b", "c"])
>>> print(td_bc)

sqrt()

计算 self 的逐元素平方根。

sqrt_()

原地版本的 sqrt()。

squeeze(*args, **kwargs)

挤压维度介于 -self.batch_dims+1 和 self.batch_dims-1 之间的所有张量，并将它们返回到一个新的 tensordict 中。

参数:

dim (Optional[int]) – 沿其挤压的维度。如果 dim 为 None，则所有单例维度都将被挤压。默认为 None。

示例

>>> td = TensorDict({
...     'x': torch.arange(24).reshape(3, 1, 4, 2),
... }, batch_size=[3, 1, 4])
>>> td = td.squeeze()
>>> td.shape
torch.Size([3, 4])
>>> td.get("x").shape
torch.Size([3, 4, 2])

此操作也可以用作上下文管理器。对原始 tensordict 的更改将发生在原地之外，即原始张量的内容不会被更改。这也假设 tensordict 未被锁定（否则，需要解锁 tensordict）。此功能与隐式挤压 不 兼容。

>>> td = TensorDict({
...     'x': torch.arange(24).reshape(3, 1, 4, 2),
... }, batch_size=[3, 1, 4])
>>> with td.squeeze(1) as tds:
...     tds.set("y", torch.zeros(3, 4))
>>> assert td.get("y").shape == [3, 1, 4]

classmethod stack(input, dim=0, *, out=None)

沿给定维度将 tensordict 堆叠成一个单一的 tensordict。

此调用等效于调用 torch.stack()，但与 torch.compile 兼容。

stack_from_tensordict(dim: int = 0, *, sorted: Optional[Union[bool, List[NestedKey]]] = None, out: Optional[Tensor] = None) → Tensor

将 tensordict 的所有条目堆叠成一个单一的张量。

参数:

dim (int, optional) – 条目应沿其堆叠的维度。

关键字参数:

• sorted (bool 或 NestedKeys 列表) – 如果为 True，则条目将按字母顺序堆叠。如果为 False（默认），将使用字典顺序。或者，可以提供键名列表，张量将相应地堆叠。这会产生一些开销，因为键列表将针对 tensordict 中的叶名称列表进行检查。

• out (torch.Tensor, optional) – 堆叠操作的可选目标张量。

stack_tensors(*keys: NestedKey, out_key: NestedKey, dim: int = 0, keep_entries: bool = False) → T

将条目堆叠到一个新条目中，并可能移除原始值。

参数:

keys (NestedKey 序列) – 要堆叠的条目。

关键字参数

out_key (NestedKey): 堆叠输入的新键名。keep_entries (bool, optional): 如果为 False，则会删除 keys 中的条目。

默认为 False。

dim (int, optional): 堆叠必须发生的维度。

默认为 0。

返回：self

示例

>>> td = TensorDict(a=torch.zeros(()), b=torch.ones(()))
>>> td.stack_tensors("a", "b", out_key="c")
>>> assert "a" not in td
>>> assert (td["c"] == torch.tensor([0, 1])).all()

state_dict(destination=None, prefix='', keep_vars=False, flatten=False) → OrderedDict[str, Any]

从 tensordict 生成 state_dict。

state_dict 的结构仍然是嵌套的，除非 flatten 设置为 True。

一个 tensordict state_dict 包含重建 tensordict 所需的所有张量和元数据（目前不支持名称）。

参数:

• destination (dict, optional) – 如果提供，tensordict 的状态将更新到 dict 中，并返回相同的对象。否则，将创建并返回一个 OrderedDict。默认值：None。

• prefix (str, optional) – 添加到张量名称的前缀，用于组成 state_dict 中的键。默认值：''。

• keep_vars (bool, optional) – 默认情况下，state_dict 中返回的 torch.Tensor 项与 autograd 分离。如果设置为 True，则不会执行分离。默认值：False。

• flatten (bool, optional) – 结构是否应使用 "." 字符展平。默认为 False。

示例

>>> data = TensorDict({"1": 1, "2": 2, "3": {"3": 3}}, [])
>>> sd = data.state_dict()
>>> print(sd)
OrderedDict([('1', tensor(1)), ('2', tensor(2)), ('3', OrderedDict([('3', tensor(3)), ('__batch_size', torch.Size([])), ('__device', None)])), ('__batch_size', torch.Size([])), ('__device', None)])
>>> sd = data.state_dict(flatten=True)
OrderedDict([('1', tensor(1)), ('2', tensor(2)), ('3.3', tensor(3)), ('__batch_size', torch.Size([])), ('__device', None)])

std(dim: Union[int, Tuple[int]] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, correction: int = 1, reduce: Optional[bool] = None) → tensordict.base.TensorDictBase | torch.Tensor

返回输入 tensordict 中所有元素的标准差值。

参数:

• dim (int, int 元组, optional) – 如果为 None，则返回一个无量纲的 tensordict，其中包含所有叶节点的总和值（如果可以计算）。如果是整数或整数元组，则仅当此维度与 tensordict 形状兼容时，才会在指定的维度上调用 std。

• keepdim (bool) – 输出张量是否保留维度。

关键字参数:

• correction (int) – 样本大小与样本自由度之间的差异。默认为贝塞尔校正，correction=1。

• reduce (bool, 可选) – 如果为 True，则将在所有 TensorDict 值上进行归约，并返回单个归约张量。默认为 False。

sub(other: tensordict.base.TensorDictBase | float, *, alpha: Optional[float] = None, default: Optional[Union[str, Tensor]] = None)

从 self 中减去 other，并按 alpha 缩放。

\[\text{{out}}_i = \text{{input}}_i - \text{{alpha}} \times \text{{other}}_i\]

支持广播、类型提升以及整数、浮点数和复数输入。

参数:

other (TensorDict, Tensor 或 Number) – 要从 self 中减去的张量或数字。

关键字参数:

• alpha (数字) – other 的乘数。

• default (torch.Tensor 或 str, 可选) – 用于独占条目的默认值。如果未提供，则两个 tensordict 的键列表必须完全匹配。如果传递了 default="intersection"，则仅考虑相交的键集，其他键将被忽略。在所有其他情况下，default 将用于操作两侧的所有缺失条目。

sub_(other: tensordict.base.TensorDictBase | float, alpha: Optional[float] = None)

原地版本的 sub()。

注意

原地 sub 不支持 default 关键字参数。

sum(dim: Union[int, Tuple[int]] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, dtype: Optional[dtype] = None, reduce: Optional[bool] = None) → tensordict.base.TensorDictBase | torch.Tensor

返回输入 tensordict 中所有元素的总和值。

参数:

• dim (int, int 元组, 可选) – 如果为 None，则返回一个无量纲的 tensordict，其中包含所有叶节点的总和值（如果可以计算）。如果为整数或整数元组，则仅当此维度与 tensordict 形状兼容时，才对指定的维度调用 sum。

• keepdim (bool) – 输出张量是否保留维度。

关键字参数:

• dtype (torch.dtype, 可选) – 返回张量的所需数据类型。如果指定，则在执行操作之前将输入张量转换为 dtype。这对于防止数据类型溢出很有用。默认值：None。

• reduce (bool, 可选) – 如果为 True，则将在所有 TensorDict 值上进行归约，并返回单个归约张量。默认为 False。

tan() → T

计算 TensorDict 中每个元素的 tan() 值。

tan_() → T

原地计算 TensorDict 中每个元素的 tan() 值。

tanh() → T

计算 TensorDict 中每个元素的 tanh() 值。

tanh_() → T

原地计算 TensorDict 中每个元素的 tanh() 值。

to(*args, **kwargs) → T

将 TensorDictBase 子类映射到另一个设备、数据类型或另一个 TensorDictBase 子类（如果允许）。

不允许将张量转换为新的数据类型，因为 tensordict 不限于包含单一张量数据类型。

参数:

• device (torch.device, optional) – tensordict 的目标设备。

• dtype (torch.dtype, optional) – tensordict 的目标浮点型或复数数据类型。

• tensor (torch.Tensor, optional) – 张量，其数据类型和设备是此 TensorDict 中所有张量的目标数据类型和设备。

关键字参数:

• non_blocking (bool, optional) – 操作是否应为非阻塞。

• memory_format (torch.memory_format, optional) – 此 tensordict 中 4D 参数和缓冲区的所需内存格式。

• batch_size (torch.Size, optional) – 输出 tensordict 的结果批大小。

• other (TensorDictBase, optional) –

  TensorDict 实例，其数据类型和设备是此 TensorDict 中所有张量的目标数据类型和设备。

  注意

  由于 TensorDictBase 实例没有数据类型，因此数据类型从示例叶子节点收集。如果存在多个数据类型，则不进行数据类型转换。

• non_blocking_pin (bool, optional) –

  如果为 True，则张量在发送到设备之前被锁定到内存。这将异步完成，但可以通过 num_threads 参数进行控制。

  注意

  调用 tensordict.pin_memory().to("cuda") 通常比 tensordict.to("cuda", non_blocking_pin=True) 慢得多，因为在第二种情况下，pin_memory 是异步调用的。如果张量很大且数量众多，则多线程 pin_memory 通常是有益的：当要发送的张量太少时，生成线程和收集数据的开销超过了多线程的优势；如果张量很小，则迭代长列表的开销也过大。

• num_threads (int or None, optional) – 如果 non_blocking_pin=True，则用于 pin_memory 的线程数。默认情况下，将生成 max(1, torch.get_num_threads()) 个线程。num_threads=0 将取消 pin_memory() 调用的任何多线程处理。

返回:

如果设备与 tensordict 设备不同，和/或传递了数据类型，则返回新的 tensordict 实例。否则返回相同的 tensordict 实例。batch_size 的修改仅在原地完成。

注意

如果 TensorDict 已合并，则结果 TensorDict 也将合并。每个新张量都将是合并存储的视图，并转换为所需设备。

示例

>>> data = TensorDict({"a": 1.0}, [], device=None)
>>> data_cuda = data.to("cuda:0")  # casts to cuda
>>> data_int = data.to(torch.int)  # casts to int
>>> data_cuda_int = data.to("cuda:0", torch.int)  # multiple casting
>>> data_cuda = data.to(torch.randn(3, device="cuda:0"))  # using an example tensor
>>> data_cuda = data.to(other=TensorDict({}, [], device="cuda:0"))  # using a tensordict example

to_dict(*, retain_none: bool = True) → dict[str, Any]

返回一个字典，其键值对与 tensordict 的键值对匹配。

参数:

retain_none (bool) – 如果为 True，则来自 tensorclass 实例的 None 值将被写入字典中。否则，它们将被丢弃。默认值：True。

to_h5(filename, **kwargs)

将 tensordict 转换为具有 h5 后端的 PersistentTensorDict。

参数:

• filename (str or path) – h5 文件的路径。

• device (torch.device or compatible, optional) – 张量返回后期望所在的设备。默认为 None（默认在 cpu 上）。

• **kwargs – 要传递给 h5py.File.create_dataset() 的 kwargs。

返回:

链接到新创建文件的 PersitentTensorDict 实例。

示例

>>> import tempfile
>>> import timeit
>>>
>>> from tensordict import TensorDict, MemoryMappedTensor
>>> td = TensorDict({
...     "a": MemoryMappedTensor.from_tensor(torch.zeros(()).expand(1_000_000)),
...     "b": {"c": MemoryMappedTensor.from_tensor(torch.zeros(()).expand(1_000_000, 3))},
... }, [1_000_000])
>>>
>>> file = tempfile.NamedTemporaryFile()
>>> td_h5 = td.to_h5(file.name, compression="gzip", compression_opts=9)
>>> print(td_h5)
PersistentTensorDict(
    fields={
        a: Tensor(shape=torch.Size([1000000]), device=cpu, dtype=torch.float32, is_shared=False),
        b: PersistentTensorDict(
            fields={
                c: Tensor(shape=torch.Size([1000000, 3]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([1000000]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([1000000]),
    device=None,
    is_shared=False)

to_module(module: Module, *, inplace: bool | None = None, return_swap: bool = True, swap_dest=None, use_state_dict: bool = False, non_blocking: bool = False, memo=None)

将 TensorDictBase 实例的内容递归地写入给定的 nn.Module 属性。

to_module 也可以用作上下文管理器，以临时使用参数/缓冲区的集合填充模块（请参见下面的示例）。

参数:

module (nn.Module) – 要将参数写入的模块。

关键字参数:

• inplace (bool, optional) – 如果为 True，则模块中的参数或张量将原地更新。默认为 False。

• return_swap (bool, optional) – 如果为 True，则将返回旧的参数配置。默认为 False。

• swap_dest (TensorDictBase, optional) – 如果 return_swap 为 True，则应写入交换的 tensordict。

• use_state_dict (bool, optional) – 如果为 True，则将使用状态字典 API 加载参数（包括状态字典钩子）。默认为 False。

• non_blocking (bool, 可选) – 如果 True 且此复制发生在不同设备之间，则复制可能会相对于主机异步发生。

示例

>>> from torch import nn
>>> module = nn.TransformerDecoder(
...     decoder_layer=nn.TransformerDecoderLayer(nhead=4, d_model=4),
...     num_layers=1)
>>> params = TensorDict.from_module(module)
>>> params.data.zero_()
>>> params.to_module(module)
>>> assert (module.layers[0].linear1.weight == 0).all()

将 tensordict 用作上下文管理器对于进行函数式调用可能很有用。 .. rubric:: 示例

>>> from tensordict import from_module
>>> module = nn.TransformerDecoder(
...     decoder_layer=nn.TransformerDecoderLayer(nhead=4, d_model=4),
...     num_layers=1)
>>> params = TensorDict.from_module(module)
>>> params = params.data * 0 # Use TensorDictParams to remake these tensors regular nn.Parameter instances
>>> with params.to_module(module):
...     # Call the module with zeroed params
...     y = module(*inputs)
>>> # The module is repopulated with its original params
>>> assert (TensorDict.from_module(module) != 0).any()

返回:

如果 return_swap 为 True，则返回包含模块值的 tensordict，否则返回 None。

to_namedtuple(dest_cls: Optional[type] = None)

将 tensordict 转换为 namedtuple。

参数:

dest_cls (Type, optional) – 可选的 namedtuple 类以使用。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> data = TensorDict({
...     "a_tensor": torch.zeros((3)),
...     "nested": {"a_tensor": torch.zeros((3)), "a_string": "zero!"}}, [3])
>>> data.to_namedtuple()
GenericDict(a_tensor=tensor([0., 0., 0.]), nested=GenericDict(a_tensor=tensor([0., 0., 0.]), a_string='zero!'))

to_padded_tensor(padding=0.0, mask_key: Optional[NestedKey] = None)

将所有嵌套张量转换为填充版本，并相应地调整批大小。

参数:

• padding (float) – tensordict 中张量的填充值。默认为 0.0。

• mask_key (NestedKey, optional) – 如果提供，则为将写入有效值掩码的键。如果异构维度不是 tensordict 批大小的一部分，则会导致错误。默认为 None

to_pytree()

将 tensordict 转换为 PyTree。

如果 tensordict 不是从 pytree 创建的，则此方法仅返回 self，而无需修改。

有关更多信息和示例，请参见 from_pytree()。

to_struct_array()

将 tensordict 转换为 numpy 结构化数组。

在 from_struct_array() - to_struct_array() 循环中，输入和输出数组的内容应匹配。但是，to_struct_array 将不会保留原始数组的内存内容。

有关更多信息，请参见 from_struct_array()。

to_tensordict(*, retain_none: Optional[bool] = None) → T

从 TensorDictBase 返回常规 TensorDict 实例。

参数:

retain_none (bool) –

如果为 True，则来自 tensorclass 实例的 None 值将被写入 tensordict 中。否则，它们将被丢弃。默认值：True。

注意

从 v0.8 开始，默认值将切换为 False。

返回:

包含相同值的新 TensorDict 对象。

transpose(dim0, dim1)

返回一个 tensordict，它是输入的转置版本。给定的维度 dim0 和 dim1 被交换。

转置 tensordict 的原地或非原地修改也会影响原始 tensordict，因为内存是共享的，并且操作被映射回原始 tensordict。

示例

>>> tensordict = TensorDict({"a": torch.randn(3, 4, 5)}, [3, 4])
>>> tensordict_transpose = tensordict.transpose(0, 1)
>>> print(tensordict_transpose.shape)
torch.Size([4, 3])
>>> tensordict_transpose.set("b",, torch.randn(4, 3))
>>> print(tensordict.get("b").shape)
torch.Size([3, 4])

trunc() → T

计算 TensorDict 中每个元素的 trunc() 值。

trunc_() → T

原地计算 TensorDict 中每个元素的 trunc() 值。

type(dst_type)

将所有张量转换为 dst_type。

参数:

dst_type (type or string) – 期望的类型

uint16()

将所有张量转换为 torch.uint16。

uint32()

将所有张量转换为 torch.uint32。

uint64()

将所有张量转换为 torch.uint64。

uint8()

将所有张量转换为 torch.uint8。

unbind(dim: int) → tuple[T, ...]

返回沿指定维度解绑的索引 tensordict 元组。

示例

>>> td = TensorDict({
...     'x': torch.arange(12).reshape(3, 4),
... }, batch_size=[3, 4])
>>> td0, td1, td2 = td.unbind(0)
>>> td0['x']
tensor([0, 1, 2, 3])
>>> td1['x']
tensor([4, 5, 6, 7])

unflatten(dim, unflattened_size)

展开 tensordict 维度，将其扩展为所需的形状。

参数:

• dim (int) – 指定要展开的输入张量的维度。

• unflattened_size (shape) – 是 tensordict 的展开维度的新形状。

示例

>>> td = TensorDict({
...     "a": torch.arange(60).view(3, 4, 5),
...     "b": torch.arange(12).view(3, 4)},
...     batch_size=[3, 4])
>>> td_flat = td.flatten(0, 1)
>>> td_unflat = td_flat.unflatten(0, [3, 4])
>>> assert (td == td_unflat).all()

unflatten_keys(separator: str = '.', inplace: bool = False) → T

递归地将扁平 tensordict 转换为嵌套 tensordict。

TensorDict 类型将丢失，结果将是一个简单的 TensorDict 实例。嵌套 tensordict 的元数据将从根推断：数据树中的所有实例将共享相同的批大小、维度名称和设备。

参数:

• separator (str, 可选) – 嵌套项之间的分隔符。

• inplace (bool, 可选) – 如果 True，则生成的 tensordict 将与进行调用的 tensordict 具有相同的标识。默认为 False。

示例

>>> data = TensorDict({"a": 1, "b - c": 2, "e - f - g": 3}, batch_size=[])
>>> data.unflatten_keys(separator=" - ")
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False),
        e: TensorDict(
            fields={
                f: TensorDict(
                    fields={
                        g: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
                    batch_size=torch.Size([]),
                    device=None,
                    is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

此方法和 unflatten_keys() 在处理状态字典时特别有用，因为它们可以无缝地将扁平字典转换为模仿模型结构的数据结构。

示例

>>> model = torch.nn.Sequential(torch.nn.Linear(3 ,4))
>>> ddp_model = torch.ao.quantization.QuantWrapper(model)
>>> state_dict = TensorDict(ddp_model.state_dict(), batch_size=[]).unflatten_keys(".")
>>> print(state_dict)
TensorDict(
    fields={
        module: TensorDict(
            fields={
                0: TensorDict(
                    fields={
                        bias: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False),
                        weight: Tensor(shape=torch.Size([4, 3]), device=cpu, dtype=torch.float32, is_shared=False)},
                    batch_size=torch.Size([]),
                    device=None,
                    is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> model_state_dict = state_dict.get("module")
>>> print(model_state_dict)
TensorDict(
    fields={
        0: TensorDict(
            fields={
                bias: Tensor(shape=torch.Size([4]), device=cpu, dtype=torch.float32, is_shared=False),
                weight: Tensor(shape=torch.Size([4, 3]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> model.load_state_dict(dict(model_state_dict.flatten_keys(".")))

unlock_() → T

解锁 tensordict 以进行非原地操作。

可以用作装饰器。

有关更多详细信息，请参见 lock_()。

unsqueeze(*args, **kwargs)

取消压缩所有维度介于 -td.batch_dims 和 td.batch_dims 之间的张量，并在新的 tensordict 中返回它们。

参数:

dim (int) – 要沿其取消压缩的维度

示例

>>> td = TensorDict({
...     'x': torch.arange(24).reshape(3, 4, 2),
... }, batch_size=[3, 4])
>>> td = td.unsqueeze(-2)
>>> td.shape
torch.Size([3, 1, 4])
>>> td.get("x").shape
torch.Size([3, 1, 4, 2])

此操作也可以用作上下文管理器。对原始 tensordict 的更改将以非原地方式发生，即原始张量的内容不会被更改。这也假设 tensordict 未被锁定（否则，需要解锁 tensordict）。

>>> td = TensorDict({
...     'x': torch.arange(24).reshape(3, 4, 2),
... }, batch_size=[3, 4])
>>> with td.unsqueeze(-2) as tds:
...     tds.set("y", torch.zeros(3, 1, 4))
>>> assert td.get("y").shape == [3, 4]

update(input_dict_or_td: T, clone: bool = False, *, keys_to_update: Optional[Sequence[NestedKey]] = None, non_blocking: bool = False, is_leaf: Optional[Callable[[Type], bool]] = None, **kwargs: Any) → T

使用来自字典或另一个 TensorDict 的值更新 TensorDict。

参数:

• input_dict_or_td (TensorDictBase or dict) – 要写入 self 的输入数据。

• clone (bool, optional) – 是否应在设置之前克隆输入（张量）字典中的张量。默认为 False。

• inplace (bool, optional) – 如果为 True 并且键与 tensordict 中的现有键匹配，则该键值对的更新将原地发生。如果找不到条目，则会添加它。默认为 False。

关键字参数:

• keys_to_update (NestedKey 序列, optional) – 如果提供，则仅更新 key_to_update 中的键列表。这旨在避免调用 data_dest.update(data_src.select(*keys_to_update))。

• non_blocking (bool, 可选) – 如果 True 且此复制发生在不同设备之间，则复制可能会相对于主机异步发生。

• is_leaf (Callable[[Type], bool], optional) – 一个可调用对象，指示对象类型是否应被视为叶子节点并交换或张量集合。

返回:

self

示例

>>> td = TensorDict({}, batch_size=[3])
>>> a = torch.randn(3)
>>> b = torch.randn(3, 4)
>>> other_td = TensorDict({"a": a, "b": b}, batch_size=[])
>>> td.update(other_td, inplace=True) # writes "a" and "b" even though they can't be found
>>> assert td['a'] is other_td['a']
>>> other_td = other_td.clone().zero_()
>>> td.update(other_td)
>>> assert td['a'] is not other_td['a']

update_(input_dict_or_td: dict[str, torch.Tensor] | tensordict.base.TensorDictBase, clone: bool = False, *, non_blocking: bool = False, **kwargs: Any) → T

在原地使用来自字典或另一个 TensorDict 的值更新 TensorDict。

与 update() 不同，如果键对于 self 是未知的，此函数将抛出错误。

参数:

• input_dict_or_td (TensorDictBase or dict) – 要写入 self 的输入数据。

• clone (bool, optional) – 是否应在设置之前克隆输入（张量）字典中的张量。默认为 False。

关键字参数:

• keys_to_update (NestedKeys 序列, 可选) – 如果提供，则仅更新 key_to_update 中的键列表。这旨在避免调用 data_dest.update_(data_src.select(*keys_to_update))。

• non_blocking (bool, 可选) – 如果 True 且此复制发生在不同设备之间，则复制可能会相对于主机异步发生。

返回:

self

示例

>>> a = torch.randn(3)
>>> b = torch.randn(3, 4)
>>> td = TensorDict({"a": a, "b": b}, batch_size=[3])
>>> other_td = TensorDict({"a": a*0, "b": b*0}, batch_size=[])
>>> td.update_(other_td)
>>> assert td['a'] is not other_td['a']
>>> assert (td['a'] == other_td['a']).all()
>>> assert (td['a'] == 0).all()

update_at_(input_dict_or_td: dict[str, torch.Tensor] | tensordict.base.TensorDictBase, index: Union[None, int, slice, str, Tensor, List[Any], Tuple[(Any, ...]], clone: bool = False, *, non_blocking: bool = False) → T

在指定索引处，使用来自字典或另一个 TensorDict 的值就地更新 TensorDict。

与 TensorDict.update 不同，如果键对于 TensorDict 是未知的，此函数将抛出错误。

参数:

• input_dict_or_td (TensorDictBase or dict) – 要写入 self 的输入数据。

• idx (int, torch.Tensor, iterable, slice) – 应该进行更新的 tensordict 的索引。

• clone (bool, 可选) – 是否应在设置之前克隆输入（tensor）字典中的张量。默认为 False。

关键字参数:

• keys_to_update (NestedKeys 序列, 可选) – 如果提供，则仅更新 key_to_update 中的键列表。

• non_blocking (bool, 可选) – 如果 True 且此复制发生在不同设备之间，则复制可能会相对于主机异步发生。

返回:

self

示例

>>> td = TensorDict({
...     'a': torch.zeros(3, 4, 5),
...     'b': torch.zeros(3, 4, 10)}, batch_size=[3, 4])
>>> td.update_at_(
...     TensorDict({
...         'a': torch.ones(1, 4, 5),
...         'b': torch.ones(1, 4, 10)}, batch_size=[1, 4]),
...    slice(1, 2))
TensorDict(
    fields={
        a: Tensor(torch.Size([3, 4, 5]), dtype=torch.float32),
        b: Tensor(torch.Size([3, 4, 10]), dtype=torch.float32)},
    batch_size=torch.Size([3, 4]),
    device=None,
    is_shared=False)
>>> assert (td[1] == 1).all()

valid_keys(include_nested: bool = False, leaves_only: bool = False, is_leaf: Optional[Callable[[Type], bool]] = None, *, sort: bool = False) → _LazyStackedTensorDictKeysView

返回 tensordict 键的生成器。

警告

TensorDict 的 keys() 方法返回键的惰性视图。如果查询了 keys 但未对其进行迭代，然后修改了 tensordict，则稍后迭代键将返回键的新配置。

参数:

• include_nested (bool, optional) – 如果 True，将返回嵌套值。默认为 False。

• leaves_only (bool, optional) – 如果 False，则仅返回叶节点。默认为 False。

• is_leaf – 一个可选的可调用对象，指示一个类是否应被视为叶节点。

关键字参数:

sort (bool, optional) – 是否应对键进行排序。对于嵌套键，键根据其连接名称排序（即，("a", "key") 将被视为 "a.key" 进行排序）。请注意，当处理大型 tensordict 时，排序可能会产生显著的开销。默认为 False。

示例

>>> from tensordict import TensorDict
>>> data = TensorDict({"0": 0, "1": {"2": 2}}, batch_size=[])
>>> data.keys()
['0', '1']
>>> list(data.keys(leaves_only=True))
['0']
>>> list(data.keys(include_nested=True, leaves_only=True))
['0', '1', ('1', '2')]

values(include_nested=False, leaves_only=False, is_leaf=None, *, sort: bool = False)

返回表示 tensordict 值的生成器。

参数:

• include_nested (bool, optional) – 如果 True，将返回嵌套值。默认为 False。

• leaves_only (bool, optional) – 如果 False，则仅返回叶节点。默认为 False。

• is_leaf – 一个可选的可调用对象，指示一个类是否应被视为叶节点。

关键字参数:

sort (bool, optional) – 是否应排序键。对于嵌套键，键根据其连接名称排序（即，("a", "key") 将被视为 "a.key" 进行排序）。请注意，当处理大型 tensordict 时，排序可能会产生显著的开销。默认为 False。

var(dim: Union[int, Tuple[int]] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, correction: int = 1, reduce: Optional[bool] = None) → tensordict.base.TensorDictBase | torch.Tensor

返回输入 tensordict 中所有元素的方差值。

参数:

• dim (int, int 元组, 可选) – 如果 None，则返回一个无量纲的 tensordict，其中包含所有叶子的总和值（如果可以计算）。如果为整数或整数元组，则仅当此维度与 tensordict 形状兼容时，才对指定的维度调用 var。

• keepdim (bool) – 输出张量是否保留维度。

关键字参数:

• correction (int) – 样本大小与样本自由度之间的差异。默认为贝塞尔校正，correction=1。

• reduce (bool, 可选) – 如果为 True，则将在所有 TensorDict 值上进行归约，并返回单个归约张量。默认为 False。

view(*shape: int, size: list | tuple | torch.Size | None = None, batch_size: torch.Size | None = None)

返回一个 tensordict，其中包含根据与 tensordict batch_size 兼容的新形状的张量的视图。

或者，可以作为第一个未命名参数提供 dtype。在这种情况下，所有张量都将使用相应的 dtype 查看。请注意，这假设新形状将与提供的 dtype 兼容。有关 dtype 视图的更多信息，请参见 view()。

参数:

• *shape (int) – 结果 tensordict 的新形状。

• dtype (torch.dtype) – 或者，用于表示张量内容的 dtype。

• size – iterable

关键字参数:

batch_size (torch.Size, 可选) – 如果提供了 dtype，则可以使用此关键字参数重置 batch-size。如果使用形状调用 view，则此参数无效。

返回:

具有所需 batch_size 的新 tensordict。

示例

>>> td = TensorDict(source={'a': torch.zeros(3,4,5),
...    'b': torch.zeros(3,4,10,1)}, batch_size=torch.Size([3, 4]))
>>> td_view = td.view(12)
>>> print(td_view.get("a").shape)  # torch.Size([12, 5])
>>> print(td_view.get("b").shape)  # torch.Size([12, 10, 1])
>>> td_view = td.view(-1, 4, 3)
>>> print(td_view.get("a").shape)  # torch.Size([1, 4, 3, 5])
>>> print(td_view.get("b").shape)  # torch.Size([1, 4, 3, 10, 1])

where(condition, other, *, out=None, pad=None)

返回一个 TensorDict，其元素根据条件从 self 或 other 中选择。

参数:

• condition (BoolTensor) – 当 True (非零) 时，产生 self，否则产生 other。

• other (TensorDictBase 或 标量) – 值（如果 other 是标量）或在条件为 False 的索引处选择的值。

关键字参数:

• out (TensorDictBase, 可选) – 输出 TensorDictBase 实例。

• pad (标量, 可选) – 如果提供，则源或目标 tensordict 中缺少的键将写为 torch.where(mask, self, pad) 或 torch.where(mask, pad, other)。默认为 None，即不容忍缺少键。

zero_() → T

就地将 tensordict 中的所有张量置零。

zero_grad(set_to_none: bool = True) → T

递归地将 TensorDict 的所有梯度置零。

参数:

set_to_none (bool, 可选) – 如果 True，则 tensor.grad 将为 None，否则为 0。默认为 True。