PersistentTensorDict

class tensordict.PersistentTensorDict(*, batch_size=None, filename=None, group=None, mode='r', backend='h5', device=None, **kwargs)

持久化 TensorDict 实现。

PersistentTensorDict 实例提供了一个与磁盘上存储的数据交互的接口，使得访问这些数据变得容易，同时仍然可以利用后端提供的快速访问优势。

与其他 TensorDictBase 子类一样，PersistentTensorDict 具有一个 device 属性。这 *并不*意味着数据存储在该设备上，而是表示加载时，数据将转换为所需的设备。

关键字参数：

• batch_size (torch.Size 或 兼容类型) – tensordict 的批次大小。默认为 torch.Size(())。

• filename (str, 可选) – h5 文件的路径。与 group 互斥。

• group (h5py.Group, 可选) – 包含数据的 h5 文件或组。与 filename 互斥。

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

• backend (str, 可选) – 存储后端。目前仅支持 "h5"。

• device (torch.device 或 兼容类型, 可选) – tensordict 的设备。默认为 None（即 PyTorch 默认设备）。

• **kwargs – 传递给 h5py.File.create_dataset() 的关键字参数。

注意

目前，当 PersistentTensorDict 实例超出作用域时不会关闭。这意味着用户有责任在必要时关闭它们。

示例

>>> import tempfile
>>> with tempfile.NamedTemporaryFile() as f:
...     data = PersistentTensorDict(file=f, batch_size=[3], mode="w")
...     data["a", "b"] = torch.randn(3, 4)
...     print(data)

abs() → T

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

abs_() → T

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

acos() → T

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

acos_() → T

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

add(other: TensorDictBase | torch.Tensor, *, alpha: float | None = None, default: str | CompatibleType | None = None) → TensorDictBase

将 other（乘以 alpha）加到 self 上。

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

参数：

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

关键字参数：

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

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

add_(other: TensorDictBase | float, *, alpha: float | None = None)

add() 的就地版本。

注意

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

addcdiv(other1: TensorDictBase | torch.Tensor, other2: 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) – 分子张量字典（或张量）

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

关键字参数：

value (数字, 可选) – \(\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) – 要乘的张量字典或张量

• other2 (TensorDict 或 Tensor) – 要乘的张量字典或张量

关键字参数：

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

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

addcmul() 的就地版本。

all(dim: int = None) → bool | TensorDictBase

检查张量字典中的所有值是否都为 True/非空。

参数：

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

any(dim: int = None) → bool | TensorDictBase

检查张量字典中是否有任何值为 True/非空。

参数：

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

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

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

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

参数：

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

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

关键字参数：

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

• device (torch.device, 可选) – 生成的设备（如果有）。

• names (字符串列表, 可选) – 新的维度名称，如果修改了 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 上执行的操作需要访问多个键才能进行单次计算，则将 out 参数设置为等于 self 可能会导致操作静默地提供错误的结果。例如

  >>> 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 构造函数的其他关键字参数。

返回值：

一个包含已转换张量的新的 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) → T

将一个可调用对象应用于 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: int | None = 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])

property batch_dims: int

TensorDict 批次大小的长度。

返回值：

表示 TensorDict 维度数量的整数。

property batch_size

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

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

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

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

返回值：

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

示例

>>> 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。

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

沿着给定维度将 TensorDict 连接成一个 TensorDict。

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

ceil() → T

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

ceil_() → T

对 TensorDict 中的每个元素计算其 ceil() 值，并进行就地操作。

chunk(chunks: int, dim: int = 0) → tuple[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: TensorDictBase | torch.Tensor, *, default: str | CompatibleType | None = None) → T

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

参数：

other (TensorDict 或 Tensor) – 另一个输入 TensorDict 或张量。

关键字参数：

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

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

clamp_max() 的就地版本。

注意

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

clamp_min(other: TensorDictBase | torch.Tensor, default: str | CompatibleType | None = None) → T

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

参数：

other (TensorDict 或 Tensor) – 另一个输入 TensorDict 或张量。

关键字参数：

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

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

clamp_min() 的就地版本。

注意

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

clear() → T

清除张量字典的内容。

clear_device_() → T

清除张量字典的设备。

返回：自身

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

将 TensorDictBase 子类实例克隆到同一类型的新的 TensorDictBase 子类上。

要从任何其他 TensorDictBase 子类型创建 TensorDict 实例，请改用 to_tensordict() 方法。

参数：

recurse (bool, 可选) – 如果为 True，则 TensorDict 中包含的每个张量也将被复制。否则，仅复制 TensorDict 树结构。默认为 True。

注意

与许多其他操作（逐点算术、形状操作等）不同，clone 不会继承原始的锁定属性。做出这种设计选择是为了能够创建可修改的克隆，这是最常见的用法。

close()

关闭持久化张量字典。

complex128()

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

complex32()

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

complex64()

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

consolidate(filename: Path | str | None = None, *, num_threads=0, device: torch.device | None = 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

将张量字典的内容整合到单个存储中，以实现快速序列化。

参数：

filename (Path, 可选) – 用于作为张量字典存储的内存映射张量的可选文件路径。

关键字参数：

• num_threads (整数, 可选) – 用于填充存储的线程数。

• device (torch.device, 可选) – 存储必须在其中实例化的可选设备。

• non_blocking (bool, 可选) – 传递给 copy_() 的 non_blocking 参数。

• inplace (bool, 可选) – 如果为 True，则结果张量字典与 self 相同，但值已更新。默认为 False。

• return_early (bool, 可选) – 如果为 True 且 num_threads>0，则该方法将返回张量字典的未来。可以使用 future.result() 查询结果张量字典。

• use_buffer (bool, 可选) – 如果为 True 且传递了文件名，则将在共享内存中创建一个中间本地缓冲区，并在最后一步将数据复制到存储位置。这可能比直接写入远程物理内存（例如 NFS）更快。默认为 False。

• share_memory (bool, 可选) – 如果为 True，则存储将放置在共享内存中。默认为 False。

• pin_memory (bool, 可选) – 整合后的数据是否应放置在固定内存中。默认为 False。

• metadata (bool, 可选) – 如果为 True，则元数据将与公共存储一起存储。如果提供了文件名，则此操作无效。当想要控制序列化方式时，存储元数据可能很有用，因为如果元数据存在或不存在，TensorDict 会以不同的方式处理整合 TD 的腌制/解腌制。

注意

如果张量字典已经整合，则所有参数都将被忽略，并返回 self。调用 contiguous() 以重新整合。

示例

>>> 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()

将PersistentTensorDict物化到一个常规的TensorDict上。

copy()

返回tensordict的浅拷贝（即复制结构但不复制数据）。

等价于 TensorDictBase.clone(recurse=False)

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

参见 TensorDictBase.update_。

non-blocking参数将被忽略，仅为了与torch.Tensor.copy_()保持兼容。

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

参见 TensorDictBase.update_at_。

cos() → T

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

cos_() → T

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

cosh() → T

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

cosh_() → T

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

cpu(**kwargs) → T

将tensordict转换为CPU。

此函数还支持to()的所有关键字参数。

create_nested(key)

创建与当前tensordict具有相同形状、设备和维度名称的嵌套tensordict。

如果值已存在，则此操作将覆盖该值。此操作在锁定的tensordict中被阻止。

示例

>>> 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

将tensordict转换为cuda设备（如果尚未在该设备上）。

参数：

device (int, 可选) – 如果提供，则为应将张量转换为的cuda设备。

此函数还支持to()的所有关键字参数。

property data

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

del_(key)

删除tensordict的一个键。

参数：

key (NestedKey) – 要删除的键

返回值：

self

property depth: int

返回tensordict的深度 - 最大层数。

最小深度为0（没有嵌套的tensordict）。

detach() → T

分离tensordict中的张量。

返回值：

一个新的tensordict，其中没有张量需要梯度。

detach_()

就地分离tensordict中的张量。

返回值：

self。

property device

TensorDict的设备。

如果TensorDict具有指定的设备，则其所有张量（包括嵌套张量）都必须位于同一设备上。如果TensorDict设备为None，则不同的值可以位于不同的设备上。

返回值：

torch.device对象，指示张量所在的设备，如果TensorDict没有设备，则为None。

示例

>>> 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

参见 batch_dims()。

div(other: TensorDictBase | torch.Tensor, *, default: str | CompatibleType | None = None) → T

将输入 self 的每个元素除以 other 的对应元素。

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

支持广播、类型提升以及整数、浮点数、TensorDict 或张量输入。始终将整数类型提升到默认标量类型。

参数：

other (TensorDict, Tensor 或 Number) – 除数。

关键字参数：

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

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

div() 的就地版本。

注意

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

double()

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

property dtype

如果 TensorDict 中值的 dtype 唯一，则返回该 dtype。

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

将 TensorDict 保存到磁盘。

此函数是 memmap() 的代理。

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

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

参数：

recurse (bool, optional) – 如果为 True，则将复制 TensorDict 的整个结构，但不包括内容。否则，仅复制根节点。默认为 False。

关键字参数：

• batch_size (torch.Size, optional) – TensorDict 的新批大小。

• device (torch.device, optional) – 新的设备。

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

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 中任一元素的张量进行就地修改会导致两个 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, **kwargs) → 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: TensorDictBase | torch.Tensor) → TensorDictBase

将张量字典的形状广播到other的形状，并相应地扩展。

如果输入是张量集合（张量字典或张量类），则叶子将一对一地扩展。

示例

>>> 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) → TensorDictBase

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

参数：

• key (str) – 要重命名的键

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

返回值：

self

filter_empty_()

就地过滤掉所有空张量字典。

filter_non_tensor_data() → T

过滤掉所有非张量数据。

flatten(start_dim=0, end_dim=- 1)

展平张量字典中的所有张量。

参数：

• 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) → T

递归地将嵌套的张量字典转换为扁平的张量字典。

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

参数：

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

• inplace (bool, optional) – 如果为True，则生成的张量字典将与调用该方法的张量字典具有相同的标识。默认为False。

• is_leaf (callable, optional) – 一个可调用的类类型，返回一个布尔值，指示该类是否必须被视为叶子。

示例

>>> 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, filename, batch_size=None, device=None, **kwargs)

将字典或 TensorDict 转换为 h5 文件。

参数：

• input_dict (dict, TensorDict 或兼容类型) – 要作为 h5 存储的数据。

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

• batch_size (tensordict 批大小, optional) – 如果提供，则为 tensordict 的批大小。否则，批大小将从输入结构（如果存在）中收集或自动确定。

• device (torch.device 或兼容类型, optional) – 返回张量后期望所在的设备。默认为None（默认在 cpu 上）。

• **kwargs – 传递给 h5py.File.create_dataset() 的关键字参数。

返回值：

一个与新创建的文件关联的PersitentTensorDict实例。

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)

类方法 from_h5(filename, mode='r')

从 h5 文件创建 PersistentTensorDict。

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

参数：

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

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

类方法 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。 .. 注意

  This is particularly useful when state-dict hooks have to be
  used.
  

示例

>>> 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)

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

通过 vmap 获取多个模块的参数，用于集成学习/特征期望应用。

参数：

modules (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。 .. 注意

  This is particularly useful when state-dict hooks have to be
  used.
  

• 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

类方法 from_namedtuple(named_tuple, *, auto_batch_size: bool = False)

递归地将命名元组转换为 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)

类方法 from_pytree(pytree, *, batch_size: torch.Size | None = None, auto_batch_size: bool = False, batch_dims: int | None = None)

将 pytree 转换为 TensorDict 实例。

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

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

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

注意

对于字典，非嵌套键（Non-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 fromkeys(keys: List[NestedKey], value: Any = 0)

根据键列表和单个值创建一个 tensordict。

参数：

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

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

gather(dim: int, index: Tensor, out: T | None = 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()。

• 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, default=_NoDefault.ZERO)

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

参数：

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

• default – 如果 tensordict 中未找到键，则为默认值。

示例

>>> td = TensorDict({"x": 1}, batch_size=[])
>>> td.get("x")
tensor(1)
>>> td.get("y", default=None)
None

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

根据键 key 和索引 idx 获取 tensordict 的值。

参数：

• key (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: NestedKey)

返回条目的形状，可能避免反复使用 get()。

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

获取非张量值（如果存在），或者如果未找到非张量值则获取 default。

此方法对张量/TensorDict 值具有鲁棒性，这意味着如果收集到的值为常规张量，它也将被返回（尽管此方法会带来一些开销，并且不应超出其自然范围）。

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

参数：

• key (NestedKey) – tensordict.tensorclass.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。

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) – 源工作进程的秩。

关键字参数：

• 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()

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

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) – 目标工作进程的秩，内容应发送到该进程。

关键字参数：

• 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: bool = False, leaves_only: bool = False, is_leaf=None) → Iterator[tuple[str, CompatibleType]]

返回 TensorDict 的键值对生成器。

参数：

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

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

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

keys(include_nested: bool = False, leaves_only: bool = False, is_leaf: Callable[[Type], bool] | None = None) → _PersistentTDKeysView

返回 TensorDict 键的生成器。

参数：

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

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

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

示例

>>> 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(input, dim=0, *, out=None, **kwargs)

创建 TensorDict 的延迟堆叠。

有关详细信息，请参见 lazy_stack()。

lerp(end: TensorDictBase | torch.Tensor, weight: 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, 张量 或 浮点数) – 插值公式的权重。

lerp_(end: TensorDictBase | float, weight: TensorDictBase | float)

lerp() 的就地版本。

lgamma() → T

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

lgamma_() → T

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

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

从磁盘加载 TensorDict。

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

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

从磁盘加载 TensorDict 到当前 TensorDict 中。

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

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

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

参数：

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

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

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

• out (TensorDictBase， 可选) – 数据应写入其中的可选 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 | Path)

将内存映射的 TensorDict 的内容加载到调用 load_memmap_ 的 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， 可选) – 是否严格强制要求 state_dict 中的键与此 TensorDict 的 torch.nn.Module.state_dict() 函数返回的键匹配。默认值：True

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

• from_flatten (bool， 可选) – 如果为 True，则假定输入状态字典已展平。默认为 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: torch.dtype | None = None) → MemoryMappedTensor

给定形状和可选的数据类型，创建一个空的内存映射张量。

警告

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

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

参数：

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

• **shape** (torch.Size 或 等效，嵌套张量的 torch.Tensor) – 要写入的张量的形状。

关键字参数：

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

返回值：

一个新的内存映射张量。

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

给定存储、形状和可选的数据类型，创建一个空的内存映射张量。

警告

此方法在设计上不是锁安全的。在多个节点上存在的内存映射 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，则新张量将共享输入的元数据（如形状和数据类型），但内容将为空。默认为 True。

返回值：

一个使用给定存储的新内存映射张量。

map(fn: Callable, dim: int = 0, num_workers: int = None, *, out: TensorDictBase = None, chunksize: int = None, num_chunks: int = None, pool: mp.Pool = 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)

将函数映射到张量字典在一个维度上的分割。

此方法将通过将张量字典分成大小相等的张量字典并在所需数量的工作进程上分发操作来将函数应用于张量字典实例。

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

参数：

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

• dim (int, 可选) – 张量字典将被分割的维度。

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

关键字参数：

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

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

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

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

• generator (torch.Generator, 可选) –

  用于播种的生成器。将从中生成一个基本种子，并且池的每个工作进程都将使用提供的种子加一个从 0 到 num_workers 的唯一整数进行播种。如果未提供生成器，则将使用随机整数作为种子。要使用未播种的工作进程，应分别创建池并直接传递到 map()。.. 注意

  Caution should be taken when providing a low-valued seed as
  this can cause autocorrelation between experiments, example:
  if 8 workers are asked and the seed is 4, the workers seed will
  range from 4 to 11. If the seed is 5, the workers seed will range
  from 5 to 12. These two experiments will have an overlap of 7
  seeds, which can have unexpected effects on the results.
  

  注意

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

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

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

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

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

• mp_start_method (str, 可选) – 多进程的启动方法。如果未提供，则将使用默认的启动方法。可接受的字符串为 "fork" 和 "spawn"。请记住，使用 "fork" 启动方法时，无法在进程之间共享 "cuda" 张量。如果将 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.]])

注意

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

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 实例。它将一次产生一个结果。

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

参数：

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

• dim (int, 可选) – 张量字典将被分割的维度。

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

关键字参数：

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

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

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

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

• generator (torch.Generator, 可选) –

  用于播种的生成器。将从中生成一个基本种子，并且池的每个工作进程都将使用提供的种子加一个从 0 到 num_workers 的唯一整数进行播种。如果未提供生成器，则将使用随机整数作为种子。要使用未播种的工作进程，应分别创建池并直接传递到 map()。.. 注意

  Caution should be taken when providing a low-valued seed as
  this can cause autocorrelation between experiments, example:
  if 8 workers are asked and the seed is 4, the workers seed will
  range from 4 to 11. If the seed is 5, the workers seed will range
  from 5 to 12. These two experiments will have an overlap of 7
  seeds, which can have unexpected effects on the results.
  

  注意

  播种工作进程的目的是让每个工作进程具有独立的种子，而不是在 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"。请记住，使用 "fork" 启动方法时，无法在进程之间共享 "cuda" 张量。如果将 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.]])

注意

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

注意

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

masked_fill(mask, value)

masked_fill 的非就地版本。

参数：

• mask (布尔型 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, value)

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

参数：

• mask (布尔型 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: TensorDictBase | torch.Tensor, *, default: str | CompatibleType | None = None) → T

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

参数：

other (TensorDict 或 Tensor) – 另一个输入 TensorDict 或张量。

关键字参数：

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

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

maximum() 的就地版本。

注意

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

classmethod maybe_dense_stack(input, dim=0, *, out=None, **kwargs)

尝试对 TensorDict 进行密集堆叠，并在需要时回退到延迟堆叠。

有关详细信息，请参阅 maybe_dense_stack()。

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

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

参数：

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

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

关键字参数：

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

• reduce (bool, 可选) – 如果为 True，则缩减将在所有 TensorDict 值上发生，并且将返回单个缩减后的张量。默认为 False。

memmap(prefix: str | None = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False) → 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，则非张量数据将在进程之间共享，并且在单个节点内的任何工作程序上进行写入操作（例如就地更新或设置）将更新所有其他工作程序上的值。如果非张量叶节点的数量很高（例如，共享大型非张量数据堆栈），这可能会导致 OOM 或类似错误。默认为 False。

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

返回值：

如果 return_early=False，则包含存储在磁盘上的张量的新的 TensorDict，否则为 TensorDictFuture 实例。

注意

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

memmap_(prefix: str | None = None, copy_existing: bool = False, num_threads: int = 0) → PersistentTensorDict

将所有张量写入相应的内存映射张量，就地执行。

参数：

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

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

关键字参数：

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

• return_early (bool, 可选) – 如果为 True 且 num_threads>0，则该方法将返回张量字典的未来。可以使用 future.result() 查询结果张量字典。

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

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

返回值：

如果return_early=False，则为self，否则为TensorDictFuture实例。

注意

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

memmap_like(prefix: str | None = None, copy_existing: bool = False, *, 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，则非张量数据将在进程之间共享，并且在单个节点内的任何工作程序上进行写入操作（例如就地更新或设置）将更新所有其他工作程序上的值。如果非张量叶节点的数量很高（例如，共享大型非张量数据堆栈），这可能会导致 OOM 或类似错误。默认为 False。

然后锁定 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: TensorDictBase | torch.Tensor, *, default: str | CompatibleType | None = None) → T

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

参数：

other (TensorDict 或 Tensor) – 另一个输入 TensorDict 或张量。

关键字参数：

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

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

minimum()的原地版本。

注意

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

mul(other: TensorDictBase | torch.Tensor, *, default: str | CompatibleType | None = 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: TensorDictBase | torch.Tensor) → T

mul()的原地版本。

注意

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

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

将一个键条件可调用函数应用于张量字典中存储的所有值，并在新的张量字典中设置它们。

可调用函数的签名必须为 Callable[Tuple[str, Tensor, ...], Optional[Union[Tensor, TensorDictBase]]]。

参数：

• fn (Callable) – 要应用于张量字典中 (名称，张量) 对的函数。对于每个叶子，只使用其叶子名称（而不是完整的 NestedKey）。

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

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

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

• device (torch.device, 可选) – 生成的设备（如果有）。

• names (字符串列表, 可选) – 新的维度名称，如果修改了 batch_size。

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

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

• filter_empty (bool, 可选) – 如果为 True，则将过滤掉空的张量字典。这还具有较低的计算成本，因为不会创建和销毁空数据结构。默认为 False，以保持向后兼容性。

• propagate_lock (bool, 可选) – 如果为 True，则锁定的张量字典将生成另一个锁定的张量字典。默认为 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 上执行的操作需要访问多个键才能进行单次计算，则将 out 参数设置为等于 self 可能会导致操作静默地提供错误的结果。例如

  >>> 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 构造函数的其他关键字参数。

返回值：

一个包含已转换张量的新的 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

张量字典的维度名称。

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

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

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

返回输入张量字典中所有非 NaN 元素的平均值。

参数：

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

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

关键字参数：

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

• reduce (bool, 可选) – 如果为 True，则缩减将在所有 TensorDict 值上发生，并且将返回单个缩减后的张量。默认为 False。

nansum(dim: int | Tuple[int] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, dtype: torch.dtype | None = None, reduce: bool | None = None) → 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

参见 batch_dims()。

ndimension() → int

参见 batch_dims()。

neg() → T

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

neg_() → T

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

new_empty(*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且包含空张量的TensorDict。

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

参数：

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

关键字参数：

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

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

• requires_grad (bool, 可选) – 是否应记录返回张量上的操作。默认值：False。

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

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

new_full(size: Size, fill_value, *, 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 (整数序列) – 定义输出张量形状的整数列表、元组或torch.Size。

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

关键字参数：

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

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

• requires_grad (bool, 可选) – 是否应记录返回张量上的操作。默认值：False。

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

• pin_memory (bool, 可选) – 如果设置，则返回的张量将在固定内存中分配。仅适用于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, 可选) – 返回的tensordict所需的类型。默认值：如果为None，则torch.dtype将保持不变。

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

• requires_grad (bool, 可选) – 是否应记录返回张量上的操作。默认值：False。

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

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

new_tensor(data: torch.Tensor | TensorDictBase, *, dtype: torch.dtype = None, device: DeviceType = _NoDefault.ZERO, requires_grad: bool = False, pin_memory: bool | None = None)

返回一个新的 TensorDict，其中 data 作为张量 data。

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

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

参数：

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

关键字参数：

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

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

• requires_grad (bool, 可选) – 是否应记录返回张量上的操作。默认值：False。

• pin_memory (bool, 可选) – 如果设置，则返回的张量将在固定内存中分配。仅适用于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, 可选) – 返回的tensordict所需的类型。默认值：如果为None，则torch.dtype将保持不变。

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

• requires_grad (bool, 可选) – 是否应记录返回张量上的操作。默认值：False。

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

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

non_tensor_items(include_nested: bool = False)

返回所有非张量叶子，可能递归。

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

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

关键字参数：

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

• dtype (torch.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!'}}

permute(*args, **kwargs)

返回根据 dims 置换批次维度的 tensordict 的视图。

参数：

• *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(*args, **kwargs)

在存储的张量上调用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 或 嵌套键) – 要查找的条目。

• 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: TensorDictBase | torch.Tensor, *, default: str | CompatibleType | None = None) → T

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

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

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

参数：

other (float, 张量 或 tensordict) – 指数值

关键字参数：

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

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

pow()的原地版本。

注意

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

prod(dim: int | Tuple[int] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, dtype: torch.dtype | None = None, reduce: bool | None = None) → 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() 值。

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

接收 tensordict 的内容并用它更新内容。

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

参数：

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

关键字参数：

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

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

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

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

在所有机器上减少 tensordict。

只有 rank 为 dst 的进程将接收最终结果。

refine_names(*names) → T

根据 names 细化 self 的维度名称。

细化是重命名的特例，它“提升”了未命名的维度。None 维度可以细化为任何名称；已命名维度只能细化为相同的名称。

因为命名张量可以与未命名张量共存，所以细化名称提供了一种编写可识别命名张量的代码的好方法，该代码适用于命名和未命名张量。

names 可以包含最多一个省略号（…）。省略号被贪婪地扩展；它被就地扩展以将 names 填充到与 self.dim() 相同的长度，使用来自 self.names 的对应索引的名称。

返回值：根据输入命名维度的相同 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) → PersistentTensorDict

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

参数：

• old_key (字符串 或 嵌套键) – 要重命名的键。

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

• safe (布尔值, 可选) – 如果为 True，则当 TensorDict 中已存在新键时抛出错误。

返回值：

self

replace(*args, **kwargs)

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

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

返回值：

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

requires_grad_(requires_grad=True) → T

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

返回此 tensordict。

参数：

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

reshape(*args, **kwargs) → T

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

参数：

*shape (int) – 结果张量字典的新形状。

返回值：

具有重新整形键的 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: str | None = None, copy_existing: bool = False, *, num_threads: int = 0, return_early: bool = False, share_non_tensor: bool = False) → T

将 TensorDict 保存到磁盘。

此函数是 memmap() 的代理。

property saved_path

返回存储内存映射保存的 TensorDict 的路径。

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

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

选择张量字典的键，并返回一个仅包含所选键的新张量字典。

值不会被复制：对原始 TensorDict 或新 TensorDict 中任一元素的张量进行就地修改会导致两个 TensorDict 都发生变化。

参数：

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

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

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

返回值：

一个新的张量字典（如果 inplace=True 则为同一个），仅包含所选键。

注意

要选择张量字典中的键并返回缺少这些键的此张量字典的版本，请参阅 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

将张量字典的内容发送到远程工作器。

参数：

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

关键字参数：

• 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, 字符串元组) – 要设置的键的名称。

• item (torch.Tensor 或等价物, TensorDictBase 实例) – 要存储在张量字典中的值。

• inplace (bool, 可选) – 如果为 True 并且键与张量字典中的现有键匹配，则将就地更新该键值对。如果 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) – 要存储在张量字典中的值

关键字参数：

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 对象。

返回：自身

示例

>>> 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

如果 key 不在 tensordict 中，则使用 default 的值为 key 条目插入。

如果 key 在 tensordict 中，则返回 key 的值，否则返回 default。

参数：

• key (str 或 嵌套键) – 值的名称。

• default (torch.Tensor 或 兼容类型, TensorDictBase) – 如果键不存在，则要存储在 tensordict 中的值。

返回值：

tensordict 中 key 的值。如果之前未设置 key，则为 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_()

将所有张量放置在共享内存中。

然后锁定 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: int | None = None) → torch.Size | int

返回由 dim 指示的维度的尺寸。

如果未指定 dim，则返回 TensorDict 的 batch_size 属性。

属性 sorted_keys: list[NestedKey]

返回按字母顺序排序的键。

不支持额外参数。

如果 TensorDict 被锁定，则键会被缓存，直到 tensordict 解锁以加快执行速度。

split(split_size: int | list[int], dim: int = 0) → list[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，可选) – 如果 True，则键将从 self 中就地删除。默认为 False。

• strict (bool，可选) – 如果 True，则在缺少键时引发异常。默认为 True。

• reproduce_struct (bool，可选) – 如果 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]

类方法 stack(input, dim=0, *, out=None)

沿给定维度将 tensordict 堆叠成单个 tensordict。

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

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，可选) – 如果提供，则 tensordict 的状态将更新到字典中并返回相同的对象。否则，将创建一个 OrderedDict 并返回。默认值：None。

• prefix (str，可选) – 添加到张量名称的用于组成 state_dict 中键的前缀。默认值：''。

• keep_vars (bool，可选) – 默认情况下，state dict 中返回的 torch.Tensor 项目与自动微分分离。如果将其设置为 True，则不会执行分离。默认值：False。

• flatten (bool，可选) – 结构是否应该使用 "." 字符进行扁平化。默认为 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: int | Tuple[int] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, correction: int = 1, reduce: bool | None = None) → TensorDictBase | torch.Tensor

返回输入 TensorDict 中所有元素的标准差值。

参数：

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

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

关键字参数：

• correction (int) – 样本大小和样本自由度之间的差异。默认为贝塞尔校正，correction=1。

• reduce (bool, 可选) – 如果为 True，则缩减将在所有 TensorDict 值上发生，并且将返回单个缩减后的张量。默认为 False。

sub(other: TensorDictBase | float, *, alpha: float | None = None, default: str | CompatibleType | None = 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: TensorDictBase | float, alpha: float | None = None)

sub() 的就地版本。

注意

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

sum(dim: int | Tuple[int] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, dtype: torch.dtype | None = None, reduce: bool | None = None) → 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: Any) → PersistentTensorDict

将 TensorDictBase 子类映射到另一个设备、数据类型或另一个 TensorDictBase 子类（如果允许）。

不允许将张量转换为新的数据类型，因为 TensorDict 不限制包含单个张量数据类型。

参数：

• device (torch.device, 可选) – TensorDict 的目标设备。

• dtype (torch.dtype, 可选) – TensorDict 的目标浮点或复数数据类型。

• tensor (torch.Tensor, 可选) – 其数据类型和设备是此 TensorDict 中所有张量的目标数据类型和设备的张量。

关键字参数：

• non_blocking (bool, 可选) – 操作是否应为非阻塞操作。

• memory_format (torch.memory_format, 可选) – 此 TensorDict 中 4D 参数和缓冲区的目标内存格式。

• batch_size (torch.Size, 可选) – 输出张量字典的结果批大小。

• other (TensorDictBase, 可选) –

  TensorDict 实例，其数据类型和设备是此 TensorDict 中所有张量的目标数据类型和设备。 .. 注意:: 由于 TensorDictBase 实例没有

  数据类型，因此数据类型是从示例叶子中收集的。如果有多个数据类型，则不进行数据类型转换。

• non_blocking_pin (bool, 可选) –

  如果 True，则在将张量发送到设备之前将其固定。这将异步执行，但可以通过 num_threads 参数进行控制。

  注意

  调用 tensordict.pin_memory().to("cuda") 通常会比 tensordict.to("cuda", non_blocking_pin=True) 慢得多，因为在第二种情况下，pin_memory 是异步调用的。如果张量很大且数量众多，多线程 pin_memory 通常会带来好处：当要发送的张量太少时，生成线程和收集数据的开销超过了多线程的益处，并且如果张量很小，遍历长列表的开销也过大。

• num_threads (int 或 None, 可选) – 如果 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() → dict[str, Any]

返回一个字典，其键值对与 tensordict 的键值对匹配。

to_h5(filename, **kwargs)

将 tensordict 转换为具有 h5 后端的 PersistentTensorDict。

参数：

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

• device (torch.device 或兼容类型, optional) – 返回张量后期望所在的设备。默认为None（默认在 cpu 上）。

• **kwargs – 传递给 h5py.File.create_dataset() 的关键字参数。

返回值：

一个与新创建的文件关联的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: nn.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 属性。

参数：

module (nn.Module) – 要将参数写入的模块。

关键字参数：

• inplace (bool, 可选) – 如果 True，则模块中的参数或张量将就地更新。默认为 False。

• return_swap (bool, 可选) – 如果 True，则将返回旧的参数配置。默认为 False。

• swap_dest (TensorDictBase, 可选) – 如果 return_swap 为 True，则为应写入交换的 tensordict。

• use_state_dict (bool, 可选) – 如果 True，则将使用 state-dict API 加载参数（包括 state-dict 钩子）。默认为 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.zero_()
>>> params.to_module(module)
>>> assert (module.layers[0].linear1.weight == 0).all()

to_namedtuple(dest_cls: type | None = None)

将 tensordict 转换为命名元组。

参数：

dest_cls (Type, 可选) – 要使用的可选命名元组类。

示例

>>> 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: NestedKey | None = None)

将所有嵌套张量转换为填充版本，并相应地调整批大小。

参数：

• padding (float) – tensordict 中张量的填充值。默认为 0.0。

• mask_key (NestedKey, 可选) – 如果提供，则为将写入有效值掩码的键。如果异构维度不是 tensordict 批大小的一部分，则会导致错误。默认为 None

to_pytree()

将 tensordict 转换为 PyTree。

如果 tensordict 不是从 pytree 创建的，则此方法仅返回 self 而不进行修改。

有关更多信息和示例，请参阅 from_pytree()。

to_tensordict() → T

从 TensorDictBase 返回一个常规的 TensorDict 实例。

返回值：

一个包含相同值的新 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 或 字符串) – 目标类型

uint16()

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

uint32()

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

uint64()

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

uint8()

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

unbind(dim: int) → tuple[T, ...]

返回一个张量字典的元组，沿着指定的维度解绑。

示例

>>> 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)

展开一个张量字典的维度，将其扩展到所需的形状。

参数：

• dim (int) – 指定要展开的输入张量的维度。

• unflattened_size (shape) – 是张量字典展开维度的新形状。

示例

>>> 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 实例。嵌套张量字典的元数据将从根节点推断：数据树中的所有实例将共享相同的批大小、维度名称和设备。

参数：

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

• inplace (bool, optional) – 如果为True，则生成的张量字典将与调用该方法的张量字典具有相同的标识。默认为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

解锁张量字典以进行非就地操作。

可用作装饰器。

有关更多详细信息，请参阅 lock_()。

unsqueeze(*args, **kwargs)

为位于 -td.batch_dims 和 td.batch_dims 之间的维度扩展所有张量，并在新的张量字典中返回它们。

参数：

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])

此操作也可以用作上下文管理器。对原始张量字典的更改将发生在外部，即原始张量的内容不会被更改。这也假设张量字典未锁定（否则，需要解锁张量字典）。

>>> 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: dict[str, CompatibleType] | T, clone: bool = False, inplace: bool = False, *, non_blocking: bool = False, keys_to_update: Sequence[NestedKey] | None = None, is_leaf: Callable[[Type], bool] | None = None) → T

使用字典或另一个张量字典中的值更新张量字典。

参数：

• input_dict_or_td (TensorDictBase 或 dict) – 要写入 self 的输入数据。

• clone (bool, 可选) – 输入（张量）字典中的张量是否应在设置之前克隆。默认为 False。

• inplace (bool, 可选) – 如果为 True 并且键与张量字典中现有键匹配，则对于该键值对，更新将在就地进行。如果找不到条目，则会添加它。默认为 False。

关键字参数：

• keys_to_update (嵌套键的序列, 可选) – 如果提供，则仅更新 key_to_update 中的键列表。这旨在避免调用 data_dest.update(data_src.select(*keys_to_update))。

• non_blocking (bool, 可选) – 如果为 True 并且此副本在不同的设备之间，则副本可能与主机异步进行。

• is_leaf (Callable[[Type], bool], 可选) – 一个可调用对象，指示对象类型是否应被视为叶子并交换或张量集合。

返回值：

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, CompatibleType] | T, clone: bool = False, *, non_blocking: bool = False, keys_to_update: Sequence[NestedKey] | None = None) → T

使用字典或另一个张量字典中的值就地更新张量字典。

与 update() 不同，如果键在 self 中未知，此函数将抛出错误。

参数：

• input_dict_or_td (TensorDictBase 或 dict) – 要写入 self 的输入数据。

• clone (bool, 可选) – 输入（张量）字典中的张量是否应在设置之前克隆。默认为 False。

关键字参数：

• keys_to_update（嵌套键的序列，可选） – 如果提供，则仅更新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, CompatibleType] | T, idx: IndexType, clone: bool = False, *, non_blocking: bool = False, keys_to_update: Sequence[NestedKey] | None = None) → T

在指定索引处就地更新 TensorDict，使用字典或另一个 TensorDict 中的值。

与 TensorDict.update 不同，如果 TensorDict 中不存在该键，此函数将抛出错误。

参数：

• input_dict_or_td (TensorDictBase 或 dict) – 要写入 self 的输入数据。

• idx（int，torch.Tensor，可迭代对象，slice） – 要进行更新的 TensorDict 的索引。

• clone（bool，可选） – 输入（张量）字典中的张量是否应在设置之前克隆。默认为 False。

关键字参数：

• keys_to_update（嵌套键的序列，可选） – 如果提供，则仅更新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()

values(include_nested: bool = False, leaves_only: bool = False, is_leaf=None) → Iterator[Tensor]

返回一个生成器，表示 TensorDict 的值。

参数：

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

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

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

var(dim: int | Tuple[int] = _NoDefault.ZERO, keepdim: bool = _NoDefault.ZERO, *, correction: int = 1, reduce: bool | None = None) → TensorDictBase | torch.Tensor

返回输入 TensorDict 中所有元素的方差值。

参数：

• dim（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 兼容。

或者，可以将数据类型作为第一个未命名的参数提供。在这种情况下，所有张量都将使用相应的数据类型进行查看。请注意，这假设新形状将与提供的数据类型兼容。有关数据类型视图的更多信息，请参阅view()。

参数：

• *shape (int) – 结果张量字典的新形状。

• dtype（torch.dtype） – 或者，用于表示张量内容的数据类型。

• size – 可迭代对象

关键字参数：

batch_size（torch.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。

参数：

• condition（BoolTensor） – 当为True（非零）时，产生self，否则产生other。

• other（TensorDictBase或标量） – 值（如果other是标量）或在 condition 为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。