TensorDictBase¶

class tensordict.TensorDictBase(*args, **kwargs)¶

TensorDictBase 是 TensorDicts 的抽象父类，TensorDicts 是 torch.Tensor 的数据容器。

abs() → Any¶: 计算 TensorDict 中每个元素的绝对值。

abs_() → Any¶: 原地计算 TensorDict 中每个元素的绝对值。

acos() → Any¶: 计算 TensorDict 中每个元素的 acos() 值。

acos_() → Any¶: 原地计算 TensorDict 中每个元素的 acos() 值。

add(other: tensordict._tensorcollection.TensorCollection | torch.Tensor, *, alpha: float | None = None, default: Optional[Union[str, Tensor, TensorCollection]] = None) → Any¶

将 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: tensordict._tensorcollection.TensorCollection | torch.Tensor | float, *, alpha: float | None = None) → Any¶: 原地版本的 add()。

注意

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

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

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

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

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

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

参数:

other1 (TensorDict 或 Tensor) – 被除数 tensordict（或 tensor）
tensor2 (TensorDict 或 Tensor) – 除数 tensordict（或 tensor）

关键字参数:

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

addcdiv_(other1, other2, *, value: float | None = 1)¶: 原地版本的 addcdiv()。

addcmul(other1: tensordict.base.TensorDictBase | torch.Tensor, other2: tensordict.base.TensorDictBase | torch.Tensor, *, value: float | None = 1) → Any¶

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

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

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

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

参数:

other1 (TensorDict 或 Tensor) – 要相乘的 tensordict 或 tensor
other2 (TensorDict 或 Tensor) – 要相乘的 tensordict 或 tensor

关键字参数:

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

addcmul_(other1, other2, *, value: float | None = 1)¶: 原地版本的 addcmul()。

abstract all(dim: Optional[int] = None) → bool | tensordict._tensorcollection.TensorCollection¶

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

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

amax(dim: int | NO_DEFAULT = NO_DEFAULT, keepdim: bool = False) → Any¶

amax(dim: int | NO_DEFAULT = NO_DEFAULT, keepdim: bool = False, *, reduce: bool) → Union[Any, Tensor]

返回输入 tensordict 中所有元素的最大值。

与 max() 相同，但 return_indices=False。

amin(dim: int | NO_DEFAULT = NO_DEFAULT, keepdim: bool = False) → Any¶

amin(dim: int | NO_DEFAULT = NO_DEFAULT, keepdim: bool = False, *, reduce: bool) → Union[Any, Tensor]

返回输入 tensordict 中所有元素的最小值。

与 min() 相同，但 return_indices=False。

abstract any(dim: Optional[int] = None) → bool | tensordict._tensorcollection.TensorCollection¶

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

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

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

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

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

参数:

fn (Callable) – 要应用于tensordict中张量的函数。
*others (TensorDictBase 实例, 可选) – 如果提供，这些 tensordict 实例应具有与 self 匹配的结构。fn 参数应接收与 tensordicts 数量（包括 self）相同的未命名输入。如果其他 tensordicts 存在缺失条目，可以通过 default 关键字参数传递默认值。

关键字参数:

batch_size (int 序列, 可选) – 如果提供，则结果 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，则会过滤掉空的 tensordicts。这也会带来更低的计算成本，因为不会创建和销毁空的数据结构。非张量数据被视为叶子节点，因此即使函数未触及它们，也会保留在 tensordict 中。默认为 False 以保持向后兼容。
propagate_lock (bool, 可选) – 如果为 True，则锁定的 tensordict 将生成另一个锁定的 tensordict。默认为 False。

call_on_nested (bool, 可选) –

如果为 True，则该函数将应用于第一级张量和容器（TensorDict 或 tensorclass）。在这种情况下，func 负责传播其调用到嵌套级别。这允许在将调用传播到嵌套 tensordicts 时进行精细控制。如果为 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) → Any¶

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

参数:

fn (Callable) – 要应用于tensordict中张量的函数。
*others (TensorDictBase 序列, 可选) – 要使用的其他tensordicts。

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

返回:: self 或 self 的副本，其中应用了函数

as_tensor() → Any¶: 将 tensordict 的每个叶子节点转换为普通的 torch.Tensor。

asin() → Any¶: 计算TensorDict中每个元素的asin()值。

asin_() → Any¶: 原地计算TensorDict中每个元素的asin()值。

atan() → Any¶: 计算TensorDict中每个元素的atan()值。

atan_() → Any¶: 原地计算TensorDict中每个元素的atan()值。

auto_batch_size_(batch_dims: Optional[int] = None, keep_compliant_size: bool = False) → Any¶

设置tensordict的最大batch-size，最多可选batch_dims。

参数:

batch_dims (int, 可选) – 如果提供，则批次大小最多为 batch_dims。
keep_compliant_size (bool, 可选) – 如果为 True，则大小符合要求的子 tensordict 在 batch_dims 被传递时，其形状不会改变。如果为 False，所有包含的 tensordicts 的 batch_dims 都将与 batch_dims 匹配。默认为 False。

返回:

self

示例

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

auto_device_() → Any¶

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

返回: self 及其编辑后的device属性。

property batch_dims: int¶

tensordict批次大小的长度。

返回:: 描述tensordict维度的整数。

abstract property batch_size: Size¶

TensorDict的形状（或batch_size）。

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

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

`batch_size`可以动态编辑，如果新大小与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() → Any¶: 将所有张量转换为torch.bfloat16。

bitwise_and(other: tensordict._tensorcollection.TensorCollection | torch.Tensor, *, default: Optional[Union[str, Tensor, TensorCollection]] = None) → Any¶

执行`self`和`other`之间的按位AND运算。

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

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

bool() → Any¶: 将所有张量转换为torch.bool。

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

计算包含的张量的字节数。

关键字参数:: count_duplicates (bool) – 是否将重复的张量计为独立的。如果为 False，则仅丢弃完全相同的张量（来自公共基张量的相同视图但不同的 ID 将被计数两次）。默认为 True（每个张量假定为单个副本）。

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

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

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

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

将tensordict的所有条目连接成一个单一的张量。

参数:

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

关键字参数:

sorted (bool 或 NestedKeys 列表) – 如果为 True，则条目将按字母顺序连接。如果为 False（默认），则使用字典顺序。或者，可以提供键名列表，并相应地连接张量。这会产生一些开销，因为键列表将与 tensordict 中的叶子名称列表进行检查。
out (torch.Tensor, 可选) – 连接操作的可选目标张量。

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

将条目连接成一个新条目，并可能删除原始值。

参数:

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

关键字参数:

out_key (NestedKey) – 连接输入的新键名。
keep_entries (bool, 可选) – 如果为 False，则 keys 中的条目将被删除。默认为 False。
dim (int, 可选) – 必须进行连接的维度。默认为 0。

返回: self

示例

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

ceil() → Any¶: 计算TensorDict中每个元素的ceil()值。

ceil_() → Any¶: 原地计算TensorDict中每个元素的ceil()值。

abstract chunk(chunks: int, dim: int = 0) → tuple[tensordict._tensorcollection.TensorCollection, ...]¶

将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(min: tensordict.base.TensorDictBase | torch.Tensor = None, max: tensordict.base.TensorDictBase | torch.Tensor = None, *, out=None) → Any¶

将 self 中的所有元素限制在 [ min, max ] 范围内。

令 min_value 和 max_value 分别为 min 和 max，返回：

\[y_i = \min(\max(x_i, \text{min\_value}_i), \text{max\_value}_i)\]

如果 min 为 None，则没有下界。或者，如果 max 为 None，则没有上界。

注意

如果 min 大于 max torch.clamp(..., min, max) 将 input 中的所有元素设置为 max 的值。

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

如果`self`中的元素大于`other`，则将它们限制为`other`。

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

clamp_max_(other: tensordict._tensorcollection.TensorCollection | torch.Tensor) → Any¶: 原地版本的 clamp_max()。

注意

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

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

将 self 的元素限制在不小于 other 的值。

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

clamp_min_(other: tensordict.base.TensorDictBase | torch.Tensor) → Any¶: 原地版本的 clamp_min()。

注意

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

clear() → Any¶: 清除 tensordict 的内容。

clear_device_() → Any¶

清除 tensordict 的设备。

返回: self

clear_refs_for_compile_() → Any¶

清除弱引用，以便 tensordict 安全地退出编译区域。

在返回 TensorDict 之前遇到 torch._dynamo.exc.Unsupported: reconstruct: WeakRefVariable() 时，请使用此方法。

返回: self

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

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

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

参数:: recurse (bool, optional) – 如果为 True，则 TensorDict 中包含的每个张量也将被复制。否则，只复制 TensorDict 的树状结构。默认为 True。

注意

与许多其他操作（逐元素算术、形状操作等）不同，clone 不会继承原始的锁定属性。此设计选择是为了能够创建一个可修改的副本，这是最常见的用法。

complex128() → Any¶: 将所有张量转换为 torch.complex128。

complex32() → Any¶: 将所有张量转换为 torch.complex32。

complex64() → Any¶: 将所有张量转换为 torch.complex64。

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

将 tensordict 内容整合到单个存储中，以实现快速序列化。

参数:

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

关键字参数:

num_threads (integer, optional) – 用于填充存储的线程数。
device (torch.device, optional) – 必须实例化存储的可选设备。
non_blocking (bool, 可选) – 传递给 copy_() 的 non_blocking 参数。
inplace (bool, 可选) – 如果为 True，则结果 tensordict 与 self 相同，但值已更新。默认为 False。
return_early (bool, 可选) – 如果为 True 且 num_threads>0，则方法将返回一个 tensordict 的 future。可以通过 future.result() 查询结果 tensordict。
use_buffer (bool, 可选) – 如果为 True 且传递了文件名，则会在共享内存中创建一个中间本地缓冲区，并将数据作为最后一步复制到存储位置。这可能比直接写入远程物理内存（例如 NFS）更快。默认为 False。
share_memory (bool, 可选) – 如果为 True，则存储将被放置在共享内存中。默认为 False。
pin_memory (bool, 可选) – 已合并数据是否应放置在固定内存中。默认为 False。
metadata (bool, 可选) – 如果为 True，则元数据将与公共存储一起存储。如果提供了文件名，则此参数无效。存储元数据在想要控制序列化方式时可能很有用，因为 TensorDict 在元数据可用或不可用时会以不同的方式处理已合并 TD 的 pickling/unpickling。

注意

如果 tensordict 已经合并，则忽略所有参数并返回 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())

abstract contiguous() → Any¶: 返回一个相同类型的新 tensordict，其值是连续的（如果值已经是连续的，则返回 self）。(Returns a new tensordict of the same type with contiguous values (or self if values are already contiguous).)

copy() → Any¶

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

等同于 TensorDictBase.clone(recurse=False)

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

请参阅 TensorDictBase.update_。

非阻塞参数将被忽略，仅为与 torch.Tensor.copy_() 的兼容性而存在。

copy_at_(tensordict: T, ) → Any¶: 请参阅 TensorDictBase.update_at_。

cos() → Any¶: 计算 TensorDict 中每个元素的 cos() 值。

cos_() → Any¶: 原地计算 TensorDict 中每个元素的 cos() 值。

cosh() → Any¶: 计算 TensorDict 中每个元素的 cosh() 值。

cosh_() → Any¶: 原地计算 TensorDict 中每个元素的 cosh() 值。

cpu(**kwargs) → Any¶

将 tensordict 转换为 CPU。

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

create_nested(key)¶

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

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

示例

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

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

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

参数:: device (int, optional) – 如果提供，则为张量应被转换到的 cuda 设备。

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

cummax(dim: int, *, return_indices: bool = True) → Any¶

cummax(dim: int, *, reduce: bool, return_indices: bool = True) → Union[Any, Tensor]

返回输入 tensordict 中所有元素的最大累积值。

参数:

dim (int) – 沿其执行 cummax 操作的维度的整数。

关键字参数:

reduce (bool, optional) – 如果为 True，则会跨所有 TensorDict 值进行归约，并返回一个单一的归约张量。默认为 False。
return_argmins (bool, optional) – 当传递 dim 参数时，cummax() 返回一个包含值和索引的命名元组。此功能的 TensorDict 等价功能是返回一个具有相同内部结构的 "values" 和 "indices" 条目的 tensorclass。默认为 True。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.cummax(dim=0)
cummax(
    indices=TensorDict(
        fields={
            a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False),
            b: TensorDict(
                fields={
                    c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.int64, is_shared=False),
                    d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False)},
                batch_size=torch.Size([4]),
                device=None,
                is_shared=False)},
        batch_size=torch.Size([4]),
        device=None,
        is_shared=False),
    vals=TensorDict(
        fields={
            a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
            b: TensorDict(
                fields={
                    c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                    d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
                batch_size=torch.Size([4]),
                device=None,
                is_shared=False)},
        batch_size=torch.Size([4]),
        device=None,
        is_shared=False),
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.cummax(reduce=True, dim=0)
torch.return_types.cummax(...)

cummin(dim: int, *, return_indices: bool = True) → Any¶

cummin(dim: int, *, reduce: bool, return_indices: bool = True) → Union[Any, Tensor]

返回输入 tensordict 中所有元素的最小累积值。

参数:

dim (int) – 沿其执行 cummin 操作的维度的整数。

关键字参数:

reduce (bool, optional) – 如果为 True，则会跨所有 TensorDict 值进行归约，并返回一个单一的归约张量。默认为 False。
return_argmins (bool, optional) – 当传递 dim 参数时，cummin() 返回一个包含值和索引的命名元组。此功能的 TensorDict 等价功能是返回一个具有相同内部结构的 "values" 和 "indices" 条目的 tensorclass。默认为 True。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.cummin(dim=0)
cummin(
    indices=TensorDict(
        fields={
            a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False),
            b: TensorDict(
                fields={
                    c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.int64, is_shared=False),
                    d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False)},
                batch_size=torch.Size([4]),
                device=None,
                is_shared=False)},
        batch_size=torch.Size([4]),
        device=None,
        is_shared=False),
    vals=TensorDict(
        fields={
            a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
            b: TensorDict(
                fields={
                    c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                    d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
                batch_size=torch.Size([4]),
                device=None,
                is_shared=False)},
        batch_size=torch.Size([4]),
        device=None,
        is_shared=False),
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.cummin(reduce=True, dim=0)
torch.return_types.cummin(...)

property data¶: 返回一个包含叶张量的 .data 属性的 tensordict。

data_ptr(*, storage: bool = False)¶

返回tensordict叶子节点的data_ptr。

这有助于检查两个tensordict是否共享相同的data_ptr()。

关键字参数:: storage (bool, optional) – 如果为 True，则会调用 tensor.untyped_storage().data_ptr()。默认为 False。

示例

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

注意

LazyStackedTensorDict 实例将显示为嵌套的 tensordicts，以反映其叶节点的实际 data_ptr()。

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

abstract del_(key: NestedKey) → Any¶

删除tensordict的键。

参数:: key (NestedKey) – 要删除的键
返回:: self

densify(layout: layout = torch.strided)¶

尝试用连续张量（普通张量或嵌套张量）来表示懒惰堆栈。

关键字参数:: layout (torch.layout) – 嵌套张量的布局，如果有的话。默认为 strided。

property depth: int¶

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

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

detach() → Any¶

分离tensordict中的张量。

返回:: 返回一个不包含需要梯度的张量的新tensordict。

abstract detach_() → Any¶

就地分离tensordict中的张量。

返回:: self。

abstract property device: torch.device | None¶

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: tensordict.base.TensorDictBase | torch.Tensor, *, default: Optional[Union[str, Tensor, TensorCollection]] = None) → Any¶

将输入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: tensordict.base.TensorDictBase | torch.Tensor) → Any¶: 原地版本的 div()。

注意

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

double() → Any¶: 将所有张量转换为torch.bool。

property dtype¶: 返回tensordict中值的dtype，如果它是唯一的。

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

将tensordict保存到磁盘。

此函数是 memmap() 的代理。

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

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

参数:

recurse (bool, optional) – 如果为 True，则将复制 TensorDict 的整个结构而没有内容。否则，仅复制根。默认为 False。

关键字参数:

batch_size (torch.Size, optional) – tensordict 的新批次大小。
device (torch.device, optional) – 新设备。
names (list of str, optional) – 维度名称。

abstract entry_class(key: NestedKey) → type¶

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

此方法应优先于 tensordict.get(key).shape，因为 get() 的执行可能很昂贵。

erf() → Any¶: 计算TensorDict中每个元素的erf()值。

erf_() → Any¶: 就地计算TensorDict中每个元素的erf()值。

erfc() → Any¶: 计算TensorDict中每个元素的erfc()值。

erfc_() → Any¶: 就地计算TensorDict中每个元素的erfc()值。

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

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

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

参数:

*keys (str) – 要排除的键。
inplace (bool) – 如果为 True，则 tensordict 将就地剪枝。默认为 False。

返回:

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

示例

>>> 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() → Any¶: 计算TensorDict中每个元素的exp()值。

exp_() → Any¶: 就地计算TensorDict中每个元素的exp()值。

abstract expand(*shape: int) → Any¶

abstract expand(shape: Size) → Any

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

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

示例

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

expand_as(other: tensordict._tensorcollection.TensorCollection | torch.Tensor) → Any¶

将tensordict的形状广播到other的形状，并相应地扩展它。

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

示例

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

expm1() → Any¶: 计算TensorDict中每个元素的expm1()值。

expm1_() → Any¶: 就地计算TensorDict中每个元素的expm1()值。

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

Fills a tensor pointed by the key with a given scalar value.

参数:

key (str or nested key) – entry to be filled.
value (Number or bool) – value to use for the filling.

返回:

self

filter_empty_()¶: 就地过滤掉所有空的tensordict。

filter_non_tensor_data() → Any¶: 过滤掉所有非张量数据。

flatten(start_dim: int | None = None, end_dim: int | None = None)¶

展平tensordict的所有张量。

参数:

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

示例

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

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

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

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

参数:

separator (str, optional) – 嵌套项之间的分隔符。
inplace (bool, optional) – 如果为 True，则结果 tensordict 将与调用它的 tensordict 具有相同的身份。默认为 False。
is_leaf (callable, optional) –
一个作用于类类型的可调用对象，返回一个布尔值，指示该类是否应被视为叶子节点。

注意

is_leaf 的目的是不是阻止递归调用嵌套的 tensordicts，而是为过滤标记某些类型，当 `leaves_only=True` 时。即使 `is_leaf(cls)` 返回 `True`，如果 `include_nested=True`，嵌套 tensordict 的结构仍将被遍历。换句话说，`is_leaf` 不控制递归深度，而是提供了一种当 `leaves_only=True` 时从结果中过滤掉某些类型的方法。这意味着树中的一个节点既可以是叶节点，也可以是具有子节点的节点。实际上，`is_leaf` 的默认值不包含 tensordict 和 tensorclass 实例作为叶节点。

另请参阅

is_leaf_nontensor() 和 default_is_leaf()。

示例

>>> 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() 特别适用于处理 state-dicts，因为它们可以无缝地将扁平字典转换为模仿模型结构的结构。

示例

>>> 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() → Any¶: 将所有张量转换为torch.float。

float16() → Any¶: 将所有张量转换为torch.float16。

float32() → Any¶: 将所有张量转换为torch.float32。

float64() → Any¶: 将所有张量转换为torch.float64。

floor() → Any¶: 计算TensorDict中每个元素的floor()值。

floor_() → Any¶: 就地计算TensorDict中每个元素的floor()值。

frac() → Any¶: 计算TensorDict中每个元素的frac()值。

frac_() → Any¶: 就地计算TensorDict中每个元素的frac()值。

classmethod from_any(obj, *, auto_batch_size: bool = False, batch_dims: Optional[int] = None, device: Optional[device] = None, batch_size: Optional[Size] = None)¶

Recursively converts any object to a TensorDict.

注意

from_any 比常规的 TensorDict 构造函数限制更少。它可以使用自定义启发式方法将数据结构（如 dataclasses 或 tuples）转换为 tensordict。此方法可能会产生一些额外的开销，并在映射策略方面涉及更多主观选择。

注意

This method recursively converts the input object to a TensorDict. If the object is already a TensorDict (or any similar tensor collection object), it will be returned as is.

参数:

obj – The object to be converted.

关键字参数:

auto_batch_size (bool, optional) – 如果为 True，则会自动计算并应用 batch size。默认为 False。
batch_dims (int, optional) – 如果 auto_batch_size 为 True，则定义输出 tensordict 应有的维度数。默认为 None (每个级别都具有完整的 batch size)。
device (torch.device, optional) – 将创建 TensorDict 的设备。
batch_size (torch.Size, optional) – TensorDict 的批次大小。与 auto_batch_size 互斥。

返回:

A TensorDict representation of the input object.

Supported objects

通过 from_dataclass() 进行数据类转换（数据类将被转换为 TensorDict 实例，而不是 tensorclasses）。
通过 from_namedtuple() 进行命名元组转换。
通过 from_dict() 进行字典转换。
通过 from_tuple() 进行元组转换。
通过 from_struct_array() 进行 NumPy 结构化数组转换。
通过 from_h5() 进行 HDF5 对象转换。

classmethod from_dataclass(dataclass, *, dest_cls: Optional[Type] = None, auto_batch_size: bool = False, batch_dims: Optional[int] = None, as_tensorclass: bool = False, device: Optional[device] = None, batch_size: Optional[Size] = None)¶

Converts a dataclass into a TensorDict instance.

参数:

dataclass – The dataclass instance to be converted.

关键字参数:

dest_cls (tensorclass, optional) – 用于映射数据的 tensorclass 类型。如果未提供，则创建一个新类。如果 `obj` 是一个类型或 as_tensorclass 为 `False`，则无效。
auto_batch_size (bool, optional) – 如果为 True，将自动确定并应用 batch size 到生成的 TensorDict。默认为 False。
batch_dims (int, optional) – 如果 auto_batch_size 为 True，则定义输出 tensordict 的维度数。默认为 None (每个级别都具有完整的 batch size)。
as_tensorclass (bool, optional) – 如果为 True，则将转换委托给自由函数 from_dataclass()，并返回一个 tensor-兼容类（tensorclass()）或实例，而不是 TensorDict。默认为 False。
device (torch.device, optional) – 将创建 TensorDict 的设备。默认为 None。
batch_size (torch.Size, optional) – TensorDict 的批次大小。默认为 None。

返回:

A TensorDict instance derived from the provided dataclass, unless as_tensorclass is True, in which case a tensor-compatible class or instance is returned.

抛出:

TypeError – If the provided input is not a dataclass instance.

警告

This method is distinct from the free function from_dataclass and serves a different purpose. While the free function returns a tensor-compatible class or instance, this method returns a TensorDict instance.

注意

此方法创建一个新的 TensorDict 实例，其键对应于输入 dataclass 的字段。
结果 TensorDict 中的每个键都使用 `cls.from_any` 方法进行初始化。
auto_batch_size 选项允许自动确定批次大小并将其应用于结果 TensorDict。

abstract classmethod from_dict(input_dict, *, auto_batch_size: Optional[bool] = None, batch_size: Optional[Size] = None, device: Optional[device] = None, batch_dims: Optional[int] = None, names: Optional[List[str]] = None)¶

从字典或另一个 TensorDict 创建 TensorDict。

If batch_size is not specified, returns the maximum batch size possible.

This function works on nested dictionaries too, or can be used to determine the batch-size of a nested tensordict.

参数:

input_dict (dictionary, optional) – a dictionary to use as a data source (nested keys compatible).

关键字参数:

auto_batch_size (bool, optional) – 如果为 True，则会自动计算并应用 batch size。默认为 False。
batch_size (iterable of int, optional) – a batch size for the tensordict.
device (torch.device 或 兼容类型, optional) – TensorDict 的设备。
batch_dims (int, optional) – batch_dims（即被视为 batch_size 的前导维度数量）。与 batch_size 互斥。请注意，这是 tensordict 的__最大__ batch dims 数量，较小的数量是可以容忍的。
names (list of str, optional) – the dimension names of the tensordict.

示例

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

abstract from_dict_instance(input_dict, *others, auto_batch_size: Optional[bool] = None, batch_size=None, device=None, batch_dims=None, names: Optional[List[str]] = None)¶

from_dict() 的实例方法版本。

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

示例

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

classmethod from_h5(filename, *, mode: str = 'r', auto_batch_size: bool = False, batch_dims: Optional[int] = None, batch_size: Optional[Size] = None)¶

从 h5 文件创建 PersistentTensorDict。

参数:

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

关键字参数:

mode (str, 可选) – 读取模式。默认为 "r"。
auto_batch_size (bool, optional) – 如果为 True，则会自动计算 batch size。默认为 False。
batch_dims (int, optional) – 如果 auto_batch_size 为 True，则定义输出 tensordict 应有的维度数。默认为 None (每个级别都具有完整的 batch size)。
batch_size (torch.Size, optional) – TensorDict 的批次大小。默认为 None。

返回:

输入 h5 文件的 PersistentTensorDict 表示。

示例

>>> td = TensorDict.from_h5("path/to/file.h5")
>>> print(td)
PersistentTensorDict(
    fields={
        key1: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False),
        key2: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

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

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

参数:

module (nn.Module) – 要从中获取参数的模块。
as_module (bool, optional) – 如果为 True，则将返回一个 TensorDictParams 实例，该实例可用于将参数存储在 torch.nn.Module 中。默认为 False。
lock (bool, optional) – 如果为 True，则结果 tensordict 将被锁定。默认为 True。
use_state_dict (bool, optional) –
如果为 True，则将使用模块的 state-dict 并将其解压到具有模型树结构的 TensorDict 中。默认为 False。

注意

这在使用 state-dict hook 时尤其有用。

示例

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

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

为 vmap 的 ensemable 学习/特征期望应用检索多个模块的参数。

参数:

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

关键字参数:

as_module (bool, optional) – 如果为 True，则会返回一个 TensorDictParams 实例，可用于将参数存储在 torch.nn.Module 中。默认为 False。
lock (bool, optional) – 如果为 True，则生成的 tensordict 将被锁定。默认为 True。
use_state_dict (bool, optional) –
如果为 True，则将使用模块的 state-dict 并将其解压到具有模型树结构的 TensorDict 中。默认为 False。

注意

这在使用 state-dict hook 时尤其有用。
lazy_stack (bool, optional) –
是否密集堆叠或懒惰堆叠参数。默认为 False（密集堆叠）。

注意

lazy_stack 和 as_module 是互斥的特性。

警告

懒惰输出和非懒惰输出之间有一个重要的区别：非懒惰输出将使用所需的批次大小重新实例化参数，而 `lazy_stack` 将仅将参数表示为懒惰堆叠。这意味着，虽然原始参数可以安全地传递给优化器（当 `lazy_stack=True` 时），但在设置为 `True` 时需要传递新参数。

警告

虽然使用懒惰堆栈来保留原始参数引用可能很诱人，但请记住，每次调用 get() 时，懒惰堆栈都会执行一次堆栈操作。这将需要内存（参数大小的 N 倍，如果构建了图，则更大）和计算时间。它还意味着优化器将包含更多参数，并且像 step() 或 zero_grad() 这样的操作将花费更长的时间来执行。总的来说，lazy_stack 应仅保留在非常少的用例中。
expand_identical (bool, optional) – 如果设置为 True 且正在将相同的参数（相同的标识）堆叠到自身，则将返回该参数的扩展版本。当 lazy_stack=True 时，将忽略此参数。

示例

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

当 lazy_stack=True 时，情况略有不同

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

classmethod from_namedtuple(named_tuple, *, auto_batch_size: bool = False, batch_dims: Optional[int] = None, device: Optional[device] = None, batch_size: Optional[Size] = None)¶

递归地将命名元组转换为 TensorDict。

参数:

named_tuple – 要转换的命名元组实例。

关键字参数:

auto_batch_size (bool, optional) – 如果为 True，则会自动计算并应用 batch size。默认为 False。
batch_dims (int, optional) – 如果 auto_batch_size 为 True，则定义输出 tensordict 的维度数。默认为 None (每个级别都具有完整的 batch size)。
device (torch.device, optional) – 将创建 TensorDict 的设备。默认为 None。
batch_size (torch.Size, optional) – TensorDict 的批次大小。默认为 None。

返回:

输入命名元组的 TensorDict 表示。

示例

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

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

将 pytree 转换为 TensorDict 实例。

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

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

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

注意

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

注意

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

示例

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

classmethod from_remote_init(src: int, group: 'ProcessGroup' | None = None, device: torch.device | None = None) → Self¶

从远程发送的元数据创建新的 tensordict 实例。

此类方法接收由 init_remote 发送的元数据，创建具有匹配形状和 dtype 的新 tensordict，然后异步接收实际的 tensordict 内容。

参数:

src (int) – 发送元数据的源进程的秩。
group ("ProcessGroup", optional) – 要使用的进程组。默认为 None。
device (torch.device, 可选) – 用于张量运算的设备。默认为 None。

返回:

使用接收到的元数据和内容初始化的新 tensordict 实例。

返回类型:

TensorDict

另请参阅

发送进程应已调用 ~.init_remote 来发送元数据和内容。

classmethod from_struct_array(struct_array: ndarray, *, auto_batch_size: bool = False, batch_dims: Optional[int] = None, device: Optional[device] = None, batch_size: Optional[Size] = None) → Any¶

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

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

注意

此方法执行零拷贝操作，这意味着生成的 TensorDict 将与输入的 numpy 数组共享相同的内存内容。因此，原地更改 numpy 数组的值会影响 TensorDict 的内容。

参数:

struct_array (np.ndarray) – 要转换的结构化 numpy 数组。

关键字参数:

auto_batch_size (bool, optional) – 如果为 True，则会自动计算 batch size。默认为 False。
batch_dims (int, optional) – 如果 auto_batch_size 为 True，则定义输出 tensordict 的维度数。默认为 None (每个级别都具有完整的 batch size)。
device (torch.device, 可选) –
将创建 TensorDict 的设备。默认为 None。

注意

更改设备（即，指定任何非 `None` 或 `"cpu"` 的设备）将传输数据，从而导致返回数据的内存位置发生更改。
batch_size (torch.Size, 可选) – TensorDict 的批次大小。默认为 None。

返回:

输入的结构化 numpy 数组的 TensorDict 表示。

示例

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

classmethod from_tuple(obj, *, auto_batch_size: bool = False, batch_dims: Optional[int] = None, device: Optional[device] = None, batch_size: Optional[Size] = None)¶

将元组转换为 TensorDict。

参数:

obj – 要转换的元组实例。

关键字参数:

auto_batch_size (bool, optional) – 如果为 True，则会自动计算 batch size。默认为 False。
batch_dims (int, optional) – 如果 auto_batch_size 为 True，则定义输出 tensordict 应有的维度数。默认为 None (每个级别都具有完整的 batch size)。
device (torch.device, optional) – 将创建 TensorDict 的设备。默认为 None。
batch_size (torch.Size, optional) – TensorDict 的批次大小。默认为 None。

返回:

输入的元组的 TensorDict 表示。

示例

>>> my_tuple = (1, 2, 3)
>>> td = TensorDict.from_tuple(my_tuple)
>>> print(td)
TensorDict(
    fields={
        0: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        1: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        2: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)

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

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

参数:

keys (list of NestedKey) – 指定新字典键的可迭代对象。
value (compatible type, optional) – 所有键的值。默认为 0。

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

沿由 dim 指定的轴收集值。

参数:

dim (int) – 要收集元素的维度
index (torch.Tensor) – 一个长整型张量，其维度数量与 tensordict 匹配，只有一个维度不同（收集维度）。它的元素引用沿所需维度收集的索引。
out (TensorDictBase, optional) – 目标 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: 'torch.distributed.ProcessGroup' | None = None) → Self | None¶

从各个工作节点收集 tensordicts 并将它们堆叠到目标节点上的 self 中。

参数:

dst (int) – 目标工作进程的 rank，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://: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://: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)¶

get(key, default)

获取输入键对应的存储值。

参数:

key (str, str 的元组) – 要查询的键。如果是 str 的元组，则等同于链式调用 getattr。
default –
如果 tensordict 中找不到该键，则返回默认值。默认为 None。

警告

以前，如果 tensordict 中不存在某个键且未提供默认值，则会引发 `KeyError`。从 v0.7 开始，此行为已更改，并返回 `None` 值（符合 dict.get 的行为）。要采用旧行为，请设置环境变量 `export TD_GET_DEFAULTS_TO_NONE=’0’` 或调用 :func`~tensordict.set_get_defaults_to_none(False)`。

注意

在处理不规则张量时，可以将关键字参数传递给 get()。有关完整概述，请参阅 get()。

示例

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

get_at(key, index)¶

get_at(key, index, default)

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

参数:

key (str, tuple of 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) – 非张量数据的存储位置。
default (Any, optional) – 找不到键时要返回的值。

返回： `tensordict.tensorclass.NonTensorData` 的内容，: 或对应于 `key` 的条目（如果它不是 `tensordict.tensorclass.NonTensorData`）或 `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() → Any¶: 将所有张量转换为 torch.half。

init_remote(dst: int, group: 'ProcessGroup' | None = None, device: torch.device | None = None) → None¶

通过发送元数据和内容来初始化远程 tensordict。

此方法将当前 tensordict 的元数据（形状、dtype 等）发送到指定的目标 rank（dst）。

然后异步发送实际的 tensordict 内容。

参数:

dst (int) – 目标进程的 rank。
group ("ProcessGroup", optional) – 要使用的进程组。默认为 None。
device (torch.device, 可选) – 用于张量运算的设备。默认为 None。

另请参阅

接收进程应调用 ~.from_remote_init 或等效方法来接收并基于发送的元数据初始化新的 tensordict。

示例

>>> import os
>>> import torch
>>> import torch.distributed as dist
>>> from tensordict import TensorDict, MemoryMappedTensor
>>> import multiprocessing as mp
>>>
>>> def server(queue):
...     # Set environment variables for distributed communication
...     os.environ["MASTER_ADDR"] = "localhost"
...     os.environ["MASTER_PORT"] = "29505"
...
...     # Initialize the distributed backend
...     dist.init_process_group("gloo", rank=0, world_size=2)
...
...     # Create a sample tensordict
...     td = (
...         TensorDict(
...             {
...                 ("a", "b"): torch.ones(2),
...                 "c": torch.ones(2),
...                 ("d", "e", "f"): MemoryMappedTensor.from_tensor(torch.ones(2, 2)),
...             },
...             [2],
...         )
...         .expand(1, 2)
...         .contiguous()
...     )
...
...     # Send the tensordict metadata and content to the client
...     td.init_remote(dst=1)
...
>>> def client(queue):
...     # Set environment variables for distributed communication
...     os.environ["MASTER_ADDR"] = "localhost"
...     os.environ["MASTER_PORT"] = "29505"
...
...     # Initialize the distributed backend
...     dist.init_process_group("gloo", rank=1, world_size=2)
...
...     # Receive the tensordict metadata and content from the server
...     received_td = TensorDict.from_remote_init(src=0)
...
...     # Verify that the received tensordict matches the expected structure and values
...     assert set(received_td.keys()) == {"a", "c", "d"}
...     assert (received_td == 1).all()
...
...     # Signal that the test has completed successfully
...     queue.put("yuppie")
>>>
>>> if __name__ == "__main__":
...     queue = mp.Queue(1)
...
...     # Create and start the server and client processes
...     main_worker = mp.Process(target=server, args=(queue,))
...     secondary_worker = mp.Process(target=client, args=(queue,))
...
...     main_worker.start()
...     secondary_worker.start()
...
...     try:
...         out = queue.get(timeout=10)  # Wait for the signal with a timeout
...         print(out)  # Should print "yuppie"
...     finally:
...         queue.close()
...         main_worker.join(timeout=10)
...         secondary_worker.join(timeout=10)

int() → Any¶: 将所有张量转换为 torch.int。

int16() → Any¶: 将所有张量转换为 torch.int16。

int32() → Any¶: 将所有张量转换为 torch.int32。

int64() → Any¶: 将所有张量转换为 torch.int64。

int8() → Any¶: 将所有张量转换为 torch.int8。

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

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

请检查 isend() 方法中的示例以了解上下文。

参数:

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

关键字参数:

group (torch.distributed.ProcessGroup, 可选) – 如果设置，将使用指定的进程组进行通信。否则，将使用默认进程组。默认为 `None`。
return_premature (bool) – 如果为 True，则返回一个 futures 列表，直到 tensordict 更新后才会返回。默认为 False，即在调用内部等待更新完成。
init_tag (int) – 用于标记张量的 init_tag。请注意，此值将按 TensorDict 中包含的张量数量进行递增。
pseudo_rand (bool) – 如果为 True，则标签序列将是伪随机的，允许来自不同节点的多数据发送而不发生重叠。请注意，这些伪随机数的生成成本很高（每秒 1e-5），这意味着它可能会减慢算法的运行速度。此值必须与传递给 isend() 的值匹配。默认为 False。

返回:

如果 return_premature=True，则返回一个未来列表以等待直到 tensordict 更新。: upon until the tensordict is updated.

is_consolidated()¶: 检查 TensorDict 是否具有合并的存储。

abstract is_contiguous() → bool¶: 返回一个布尔值，指示所有张量是否是连续的。

is_empty() → bool¶: 检查 tensordict 是否包含任何叶子。

is_floating_point() → 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: 'torch.distributed.ProcessGroup' | None = None, init_tag: int = 0, pseudo_rand: bool = False, return_early: bool = False) → int | List['Work']¶

异步发送 tensordict 的内容。

参数:

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

关键字参数:

group (torch.distributed.ProcessGroup, 可选) – 如果设置，将使用指定的进程组进行通信。否则，将使用默认进程组。默认为 `None`。
init_tag (int) – 用于标记张量的初始标签。注意，该标签将按 TensorDict 中包含的张量数量进行递增。
pseudo_rand (bool) – 如果为 True，则标签序列将是伪随机的，允许来自不同节点的多数据发送而不发生重叠。请注意，这些伪随机数的生成成本很高（每秒 1e-5），这意味着它可能会减慢算法的运行速度。默认为 False。
return_early (bool, optional) – 如果为 True，则返回一个 futures 列表，而不是最后一个已发送张量的标签。默认为 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://: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://: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() → Any¶

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

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

isnan() → Any¶

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

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

isneginf() → Any¶: 测试输入中的每个元素是否为负无穷。

isposinf() → Any¶: 测试输入中的每个元素是否为负无穷。

isreal() → Any¶: 返回一个新的 tensordict，其中包含布尔值元素，表示输入中的每个元素是否为实值。

items(include_nested: bool = False, leaves_only: bool = False, is_leaf=None, *, sort: bool = False) → Iterator[tuple[str, Union[torch.Tensor, tensordict._tensorcollection.TensorCollection]]]¶

返回 tensordict 的键值对生成器。

参数:

include_nested (bool, 可选) – 如果为 `True`，则返回嵌套值。默认为 `False`。
leaves_only (bool, 可选) – 如果为 `False`，则仅返回叶子节点。默认为 `False`。
is_leaf (callable, optional) –
一个作用于类类型的可调用对象，返回一个布尔值，指示该类是否应被视为叶子节点。

注意

is_leaf 的目的是不是阻止递归调用嵌套的 tensordicts，而是为过滤标记某些类型，当 `leaves_only=True` 时。即使 `is_leaf(cls)` 返回 `True`，如果 `include_nested=True`，嵌套 tensordict 的结构仍将被遍历。换句话说，`is_leaf` 不控制递归深度，而是提供了一种当 `leaves_only=True` 时从结果中过滤掉某些类型的方法。这意味着树中的一个节点既可以是叶节点，也可以是具有子节点的节点。实际上，`is_leaf` 的默认值不包含 tensordict 和 tensorclass 实例作为叶节点。

另请参阅

is_leaf_nontensor() 和 default_is_leaf()。

关键字参数:

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

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

返回 tensordict 键的生成器。

警告

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

参数:

include_nested (bool, 可选) – 如果为 `True`，则返回嵌套值。默认为 `False`。
leaves_only (bool, 可选) – 如果为 `False`，则仅返回叶子节点。默认为 `False`。
is_leaf (callable, optional) –
一个作用于类类型的可调用对象，返回一个布尔值，指示该类是否应被视为叶子节点。

注意

is_leaf 的目的是不是阻止递归调用嵌套的 tensordicts，而是为过滤标记某些类型，当 `leaves_only=True` 时。即使 `is_leaf(cls)` 返回 `True`，如果 `include_nested=True`，嵌套 tensordict 的结构仍将被遍历。换句话说，`is_leaf` 不控制递归深度，而是提供了一种当 `leaves_only=True` 时从结果中过滤掉某些类型的方法。这意味着树中的一个节点既可以是叶节点，也可以是具有子节点的节点。实际上，`is_leaf` 的默认值不包含 tensordict 和 tensorclass 实例作为叶节点。

另请参阅

is_leaf_nontensor() 和 default_is_leaf()。

关键字参数:

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

示例

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

classmethod lazy_stack(input, dim: int = 0, *, out=None, **kwargs)¶

创建 TensorDicts 的懒惰堆叠。

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

lerp(end: tensordict._tensorcollection.TensorCollection | torch.Tensor, weight: tensordict._tensorcollection.TensorCollection | torch.Tensor | float) → Any¶

根据标量或张量 `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, 张量或float) – 插值公式的权重。

lerp_(end: tensordict.base.TensorDictBase | torch.Tensor | float, weight: tensordict.base.TensorDictBase | torch.Tensor | float)¶: 原地版本的 lerp()。

lgamma() → Any¶: 计算 TensorDict 中每个元素的 lgamma() 值。

lgamma_() → Any¶: 原地计算 TensorDict 中每个元素的 lgamma() 值。

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

从磁盘加载 tensordict。

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

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

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

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

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

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

参数:

prefix (str 或 文件夹路径) – 应从中获取已保存 tensordict 的文件夹路径。
device (torch.device 或 等效项, 可选) – 如果提供，数据将异步转换为该设备。支持 `"meta"` 设备，在这种情况下，数据不会被加载，而是创建一组空的 "meta" 张量。这对于在不实际打开任何文件的情况下了解模型大小和结构很有用。
non_blocking (bool, 可选) – 如果为 `True`，则在将张量加载到设备后不会调用同步。默认为 `False`。
out (TensorDictBase, optional) – 可选的 tensordict，数据应写入其中。

示例

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

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

示例

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

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

示例

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

load_memmap_(prefix: str | pathlib.Path)¶

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

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

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

将 state-dict（格式如 state_dict()）加载到 tensordict 中。

参数:

state_dict (OrderedDict) – 要复制的 state_dict。
strict (bool, optional) – 是否严格要求 state_dict 中的键与此 tensordict 的 torch.nn.Module.state_dict() 函数返回的键匹配。默认值：True
assign (bool, optional) – 是否将 state 字典中的项分配给 tensordict 中相应的键，而不是就地复制到 tensordict 的当前张量中。当 False 时，将保留当前模块张量的属性；当 True 时，将保留 state 字典中张量的属性。默认为 False。
from_flatten (bool, optional) – 如果为 True，则假定输入的 state_dict 是扁平化的。默认为 False。

示例

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

lock_() → Any¶

锁定 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() → Any¶: 计算 TensorDict 中每个元素的 log() 值。

log10() → Any¶: 计算 TensorDict 中每个元素的 log10() 值。

log10_() → Any¶: 原地计算 TensorDict 中每个元素的 log10() 值。

log1p() → Any¶: 计算 TensorDict 每个元素的 log1p() 值。

log1p_() → Any¶: 原地计算 TensorDict 每个元素的 log1p() 值。

log2() → Any¶: 计算 TensorDict 每个元素的 log2() 值。

log2_() → Any¶: 原地计算 TensorDict 每个元素的 log2() 值。

log_() → Any¶: 原地计算 TensorDict 每个元素的 log() 值。

logical_and(other: tensordict._tensorcollection.TensorCollection | torch.Tensor, *, default: Optional[Union[str, Tensor, TensorCollection]] = None) → Any¶

对 self 和 other 执行逻辑 AND 操作。

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

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

logsumexp(dim=None, keepdim=False, *, out=None)¶

返回给定维度 dim 上输入 tensordict 各行的指数对数之和。计算是数值稳定的。

如果 keepdim 为 True，则输出张量的大小与输入张量相同，但在 dim 维度上大小为 1。否则，dim 将被压缩（参见 squeeze()），导致输出张量比输入张量少 1 个（或 len(dim) 个）维度。

参数:

dim (int or tuple of ints, optional) – 要约简的维度或维度。如果为 None，则约简 tensordict 的所有批次维度。
keepdim (bool) – 输出 tensordict 是否保留维度。

关键字参数:

out (TensorDictBase, optional) – 输出的 tensordict。

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

根据形状和可选的 dtype 创建一个空的内存映射张量。

警告

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

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

参数:

key (NestedKey) – 要写入的新条目的键。如果键已存在于 tensordict 中，将引发异常。
shape (torch.Size or equivalent, torch.Tensor for nested tensors) – 要写入的张量的形状。

关键字参数:

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

返回:

一个新的内存映射张量。

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

根据存储、形状和可选的 dtype 创建一个空的内存映射张量。

警告

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

注意

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

参数:

key (NestedKey) – 要写入的新条目的键。如果键已存在于 tensordict 中，将引发异常。
storage (torch.UntypedStorage) – 用于新 MemoryMappedTensor 的存储。必须是物理内存存储。
shape (torch.Size or equivalent, torch.Tensor for nested tensors) – 要写入的张量的形状。

关键字参数:

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

返回:

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

abstract 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, optionaL) – 如果为 False，则新张量将共享输入张量的元数据（如形状和 dtype），但内容将为空。默认为 True。

返回:

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

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

将一个函数映射到tensordict在某个维度上的切片。

此方法将通过将 tensordict 切块成相等大小的 tensordict 并分派操作到所需的 Worker 数量来应用一个函数。

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

注意

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

参数:

fn (callable) – 要应用于 tensordict 的函数。支持类似于 Callabe[[TensorDict], Union[TensorDict, Tensor]] 的签名。
dim (int, 可选) – tensordict 将被分块的维度。
num_workers (int, optional) – 工作线程的数量。与 pool 互斥。如果未提供，工作线程数量将设置为可用 CPU 的数量。

关键字参数:

out (TensorDictBase, optional) – 一个可选的输出容器。沿提供的 dim 的批次大小必须与 self.ndim 匹配。如果它是共享的或 memmap 的（is_shared() 或 is_memmap() 返回 True），它将在远程进程中被填充，避免数据向内传输。否则，来自 self 切片的将发送到进程，在当前进程中收集，然后原地写入 out。
chunksize (int, optional) – 每个数据块的大小。chunksize 为 0 将沿期望的维度解绑 tensordict 并在函数应用后重新堆叠它，而 chunksize>0 将分割 tensordict 并对生成的 tensordict 列表调用 torch.cat()。如果未提供，块的数量将等于 worker 的数量。对于非常大的 tensordicts，如此大的块可能无法在内存中容纳以完成操作，可能需要更多的块才能使操作实际上可行。此参数与 num_chunks 互斥。
num_chunks (int, optional) – 将 tensordict 分割成的块的数量。如果未提供，块的数量将等于 worker 的数量。对于非常大的 tensordicts，如此大的块可能无法在内存中容纳以完成操作，可能需要更多的块才能使操作实际上可行。此参数与 chunksize 互斥。
pool (mp.Pool, optional) – 要用于执行作业的多进程池实例。如果未提供，将在 map 方法中创建一个池。
generator (torch.Generator, optional) –
用于播种的生成器。将从中生成一个基础种子，并且池中的每个 worker 将使用提供的种子（加上从 0 到 num_workers 的唯一整数）进行播种。如果未提供生成器，将使用一个随机整数作为种子。要使用未播种的 worker，应单独创建一个池并将其直接传递给 map()。

注意

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

注意

为工作线程设置种子的目的是让每个工作线程拥有不同的种子，而不是为了让 map 方法的调用结果可重现。换句话说，两个实验可能并且很可能返回不同的结果，因为无法知道哪个工作线程会执行哪个作业。但是，我们可以确保每个工作线程都有不同的种子，并且每个工作线程上的伪随机操作将是不相关的。
max_tasks_per_child (int, optional) – 每个子进程拾取的最大任务数。默认为 None，即对任务数没有限制。
worker_threads (int, optional) – worker 的线程数。默认为 1。
index_with_generator (bool, optional) – 如果为 True，则在查询期间完成 tensordict 的分割/分块，从而节省初始化时间。请注意，chunk() 和 split() 比索引（在生成器中使用）效率更高，因此初始化时间的处理时间增加可能会对总运行时间产生负面影响。默认为 False。
pbar (bool, optional) – 如果为 True，将显示进度条。需要安装 tqdm。默认为 False。
mp_start_method (str, optional) – 多进程的启动方法。如果未提供，将使用默认启动方法。可接受的字符串是 "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.]])

map_iter(fn: Callable[[TensorCollection], TensorCollection | 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) → Iterator[T]¶

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

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

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

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

注意

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

注意

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

参数:

fn (callable) – 要应用于 tensordict 的函数。支持类似于 Callabe[[TensorDict], Union[TensorDict, Tensor]] 的签名。
dim (int, 可选) – tensordict 将被分块的维度。
num_workers (int, optional) – 工作线程的数量。与 pool 互斥。如果未提供，工作线程数量将设置为可用 CPU 的数量。

关键字参数:

shuffle (bool, optional) – 索引是否应该被全局打乱。如果为 True，每个批次将包含不连续的样本。如果 index_with_generator=False 且 shuffle=True`，则会引发错误。默认为 False。
chunksize (int, optional) – 每个数据块的大小。chunksize 为 0 将沿期望的维度解绑 tensordict 并在函数应用后重新堆叠它，而 chunksize>0 将分割 tensordict 并对生成的 tensordict 列表调用 torch.cat()。如果未提供，块的数量将等于 worker 的数量。对于非常大的 tensordicts，如此大的块可能无法在内存中容纳以完成操作，可能需要更多的块才能使操作实际上可行。此参数与 num_chunks 互斥。
num_chunks (int, optional) – 将 tensordict 分割成的块的数量。如果未提供，块的数量将等于 worker 的数量。对于非常大的 tensordicts，如此大的块可能无法在内存中容纳以完成操作，可能需要更多的块才能使操作实际上可行。此参数与 chunksize 互斥。
pool (mp.Pool, optional) – 要用于执行作业的多进程池实例。如果未提供，将在 map 方法中创建一个池。
generator (torch.Generator, optional) –
用于播种的生成器。将从中生成一个基础种子，并且池中的每个 worker 将使用提供的种子（加上从 0 到 num_workers 的唯一整数）进行播种。如果未提供生成器，将使用一个随机整数作为种子。要使用未播种的 worker，应单独创建一个池并将其直接传递给 map()。

注意

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

注意

为工作线程设置种子的目的是让每个工作线程拥有不同的种子，而不是为了让 map 方法的调用结果可重现。换句话说，两个实验可能并且很可能返回不同的结果，因为无法知道哪个工作线程会执行哪个作业。但是，我们可以确保每个工作线程都有不同的种子，并且每个工作线程上的伪随机操作将是不相关的。
max_tasks_per_child (int, optional) – 每个子进程拾取的最大任务数。默认为 None，即对任务数没有限制。
worker_threads (int, optional) – worker 的线程数。默认为 1。
index_with_generator (bool, optional) –
如果为 True，则在查询期间完成 tensordict 的分割/分块，从而节省初始化时间。请注意，chunk() 和 split() 比索引（在生成器中使用）效率更高，因此初始化时间的处理时间增加可能会对总运行时间产生负面影响。默认为 True。

注意

index_with_generator 的默认值对于 map_iter 和 map 是不同的，前者假定将拆分后的 TensorDict 存储在内存中成本过高。
pbar (bool, optional) – 如果为 True，将显示进度条。需要安装 tqdm。默认为 False。
mp_start_method (str, optional) – 多进程的启动方法。如果未提供，将使用默认启动方法。可接受的字符串是 "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.]])

abstract masked_fill(mask: Tensor, value: float | bool) → Any¶

masked_fill 的非原地版本。

参数:

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

返回:

self

示例

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

abstract masked_fill_(mask: Tensor, value: float | bool) → Any¶

用期望值填充与掩码对应的项。

参数:

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

返回:

self

示例

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

abstract masked_select(mask: Tensor) → Any¶

屏蔽 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.]])

max(dim: int | NO_DEFAULT = NO_DEFAULT, keepdim: bool = False, *, return_indices: bool = True) → Any¶

max(dim: int | NO_DEFAULT = NO_DEFAULT, keepdim: bool = False, *, reduce: bool, return_indices: bool = True) → Union[Any, Tensor]

返回输入 tensordict 中所有元素的最大值。

参数:

dim (int, optional) – 如果为 None，则返回一个无维度的 tensordict，其中包含所有叶子的最大值（如果可以计算）。如果为整数，则仅当该维度与 tensordict 的形状兼容时，才会在指定的维度上调用 max。
keepdim (bool) – 输出张量是否保留维度。

关键字参数:

reduce (bool, optional) – 如果为 True，则会跨所有 TensorDict 值进行归约，并返回一个单一的归约张量。默认为 False。
return_argmins (bool, optional) – max() 当传递 dim 参数时，返回一个包含值和索引的命名元组。这等同于 TensorDict 返回一个具有相同结构但条目为 "values" 和 "indices" 的 tensordict。默认为 True。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.max(dim=0)
max(
    indices=TensorDict(
        fields={
            a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False),
            b: TensorDict(
                fields={
                    c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.int64, is_shared=False),
                    d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False)},
                batch_size=torch.Size([4]),
                device=None,
                is_shared=False)},
        batch_size=torch.Size([4]),
        device=None,
        is_shared=False),
    vals=TensorDict(
        fields={
            a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
            b: TensorDict(
                fields={
                    c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                    d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
                batch_size=torch.Size([4]),
                device=None,
                is_shared=False)},
        batch_size=torch.Size([4]),
        device=None,
        is_shared=False),
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td.max()
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([]), 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)
>>> td.max(reduce=True)
tensor(3.2942)

maximum(other: tensordict._tensorcollection.TensorCollection | torch.Tensor, *, default: Optional[Union[str, Tensor, TensorCollection]] = None) → Any¶

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

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

maximum_(other: tensordict._tensorcollection.TensorCollection | torch.Tensor) → Any¶: 原地版本的 maximum()。

注意

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

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

尝试使 TensorDicts 密集堆叠，并在需要时回退到懒惰堆叠。

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

mean(dim: Union[int, Tuple[int], Literal['feature']] = NO_DEFAULT, keepdim: bool = NO_DEFAULT, *, dtype: torch.dtype | None = None) → Any¶

mean(dim: Union[int, Tuple[int], Literal['feature']] = NO_DEFAULT, keepdim: bool = NO_DEFAULT, *, dtype: torch.dtype | None = None, reduce: bool) → Union[Any, Tensor]

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

参数:

dim (int, tuple of int, str, optional) – 如果为 None，则返回一个无维度的 tensordict，包含所有叶子的平均值（如果可计算）。如果为整数或整数元组，则在指定的维度上调用 mean，前提是该维度与 tensordict 的形状兼容。目前只允许 “feature” 字符串。使用 dim=”feature” 将在所有特征维度上进行约简。如果 reduce=True，将返回一个形状与 TensorDict 的批次大小相同的张量。否则，将返回一个与 self 具有相同结构的、经过特征维度约简的新 tensordict。
keepdim (bool) – 输出张量是否保留维度。

关键字参数:

dtype (torch.dtype, optional) – 返回张量所需的 dtype。如果指定，则在执行操作之前将输入张量转换为 dtype。这有助于防止数据类型溢出。默认：None。
reduce (bool, optional) – 如果为 True，则会跨所有 TensorDict 值进行归约，并返回一个单一的归约张量。默认为 False。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.mean(dim=0)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td.mean()
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([]), 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)
>>> td.mean(reduce=True)
tensor(-0.0547)
>>> td.mean(dim="feature")
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3, 4]),
    device=None,
    is_shared=False)
>>> td = TensorDict(
...     a=torch.ones(3, 4, 5),
...     b=TensorDict(
...         c=torch.ones(3, 4, 5),
...         d=torch.ones(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.mean(reduce=True, dim="feature")
tensor([[1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.]])
>>> td.mean(reduce=True, dim=0)
tensor([[1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.]])

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

将所有张量写入内存映射的 Tensor 中，并放入新的 tensordict。

参数:

prefix (str) – 内存映射张量将存储的目录前缀。目录树结构将模仿 tensordict 的结构。
copy_existing (bool) – 如果为 False（默认值），并且 tensordict 中某项已是存储在磁盘上的张量且关联了文件，但未按 prefix 保存到正确位置，则会引发异常。如果为 True，则任何现有张量都将被复制到新位置。

关键字参数:

num_threads (int, optional) – 用于写入 memmap 张量的线程数。默认为 0。
return_early (bool, optional) – 如果为 True 且 num_threads>0，则该方法将返回一个 tensordict 的 future。
share_non_tensor (bool, optional) – 如果为 True，非张量数据将在进程之间共享，并且在一个节点内的任何 worker 上的写入操作（如原地更新或设置）将更新所有其他 worker 上的值。如果非张量叶子的数量很高（例如，共享大量非张量数据），这可能会导致 OOM 或类似错误。默认为 False。
existsok (bool, optional) – 如果为 False，如果同一路径中已存在张量，则会引发异常。默认为 True。

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

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

注意

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

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

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

参数:

prefix (str) – 内存映射张量将存储的目录前缀。目录树结构将模仿 tensordict 的结构。
copy_existing (bool) – 如果为 False（默认值），并且 tensordict 中某项已是存储在磁盘上的张量且关联了文件，但未按 prefix 保存到正确位置，则会引发异常。如果为 True，则任何现有张量都将被复制到新位置。

关键字参数:

num_threads (int, optional) – 用于写入 memmap 张量的线程数。默认为 0。
return_early (bool, 可选) – 如果为 True 且 num_threads>0，则方法将返回一个 tensordict 的 future。可以通过 future.result() 查询结果 tensordict。
share_non_tensor (bool, optional) – 如果为 True，非张量数据将在进程之间共享，并且在一个节点内的任何 worker 上的写入操作（如原地更新或设置）将更新所有其他 worker 上的值。如果非张量叶子的数量很高（例如，共享大量非张量数据），这可能会导致 OOM 或类似错误。默认为 False。
existsok (bool, optional) – 如果为 False，如果同一路径中已存在张量，则会引发异常。默认为 True。

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

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

注意

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

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

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

参数:

prefix (str) – 内存映射张量将存储的目录前缀。目录树结构将模仿 tensordict 的结构。
copy_existing (bool) – 如果为 False（默认值），并且 tensordict 中某项已是存储在磁盘上的张量且关联了文件，但未按 prefix 保存到正确位置，则会引发异常。如果为 True，则任何现有张量都将被复制到新位置。

关键字参数:

num_threads (int, optional) – 用于写入 memmap 张量的线程数。默认为 0。
return_early (bool, optional) – 如果为 True 且 num_threads>0，则该方法将返回一个 tensordict 的 future。
share_non_tensor (bool, optional) – 如果为 True，非张量数据将在进程之间共享，并且在一个节点内的任何 worker 上的写入操作（如原地更新或设置）将更新所有其他 worker 上的值。如果非张量叶子的数量很高（例如，共享大量非张量数据），这可能会导致 OOM 或类似错误。默认为 False。
existsok (bool, optional) – 如果为 False，如果同一路径中已存在张量，则会引发异常。默认为 True。

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

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

注意

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

示例

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

memmap_refresh_()¶

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

如果没有任何路径与之关联，此方法将引发异常。

min(dim: int | NO_DEFAULT = NO_DEFAULT, keepdim: bool = False, *, return_indices: bool = True) → Any¶

min(dim: int | NO_DEFAULT = NO_DEFAULT, keepdim: bool = False, *, reduce: bool, return_indices: bool = True) → Union[Any, Tensor]

返回输入 tensordict 中所有元素的最小值。

参数:

dim (int, optional) – 如果 None，则返回一个无维度的 tensordict，其中包含所有叶子的最小值（如果可以计算）。如果为整数，则仅当此维度与 tensordict 的形状兼容时，才会对指定的维度调用 min。
keepdim (bool) – 输出张量是否保留维度。

关键字参数:

reduce (bool, optional) – 如果为 True，则会跨所有 TensorDict 值进行归约，并返回一个单一的归约张量。默认为 False。
return_argmins (bool, optional) – min() 当传递 dim 参数时，返回一个包含值和索引的命名元组。这等同于 TensorDict 返回一个具有相同结构但条目为 "values" 和 "indices" 的 tensordict。默认为 True。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.min(dim=0)
min(
    indices=TensorDict(
        fields={
            a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False),
            b: TensorDict(
                fields={
                    c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.int64, is_shared=False),
                    d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.int64, is_shared=False)},
                batch_size=torch.Size([4]),
                device=None,
                is_shared=False)},
        batch_size=torch.Size([4]),
        device=None,
        is_shared=False),
    vals=TensorDict(
        fields={
            a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
            b: TensorDict(
                fields={
                    c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                    d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
                batch_size=torch.Size([4]),
                device=None,
                is_shared=False)},
        batch_size=torch.Size([4]),
        device=None,
        is_shared=False),
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td.min()
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([]), 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)
>>> td.min(reduce=True)
tensor(-2.9953)

minimum(other: tensordict._tensorcollection.TensorCollection | torch.Tensor, *, default: Optional[Union[str, Tensor, TensorCollection]] = None) → Any¶

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

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

minimum_(other: tensordict._tensorcollection.TensorCollection | torch.Tensor) → Any¶: 原地版本的 minimum()。

注意

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

mul(other: tensordict._tensorcollection.TensorCollection | torch.Tensor, *, default: Optional[Union[str, Tensor, TensorCollection]] = None) → Any¶

将 other 乘以 self。

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

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

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

mul_(other: tensordict._tensorcollection.TensorCollection | torch.Tensor) → Any¶: 原地版本的 mul()。

注意

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

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

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

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

参数:

fn (Callable) – 要应用于 tensordict 中的（名称，张量）对的函数。对于每个叶子，只使用其叶子名称（不使用完整的 NestedKey）。
*others (TensorDictBase 实例, 可选) – 如果提供，这些 tensordict 实例应具有与 self 匹配的结构。fn 参数应接收与 tensordicts 数量（包括 self）相同的未命名输入。如果其他 tensordicts 存在缺失条目，可以通过 default 关键字参数传递默认值。
nested_keys (bool, optional) – 如果为 True，则将使用叶子的完整路径。默认为 False，即仅将最后一个字符串传递给函数。
batch_size (int 序列, 可选) – 如果提供，则结果 TensorDict 将具有所需的 batch_size。batch_size 参数应与转换后的 batch_size 匹配。这是一个仅关键字参数。
device (torch.device, 可选) – 生成的设备，如果存在。
names (字符串列表, 可选) – 新的维度名称，以防batch_size被修改。
inplace (bool, optional) – 如果为 True，则就地进行更改。默认为 False。这是一个仅关键字参数。
default (Any, 可选) – 其他tensordict中缺失条目的默认值。如果未提供，缺失的条目将引发KeyError。
filter_empty (bool, optional) – 如果为 True，则将过滤掉空的 tensordicts。这还将降低计算成本，因为不会创建和销毁空的数据结构。为了向后兼容，默认为 False。
propagate_lock (bool, optional) – 如果为 True，则锁定的 tensordict 将产生另一个锁定的 tensordict。默认为 False。

call_on_nested (bool, optional) –

如果为 True，则该函数将应用于第一级张量和容器（TensorDict 或 tensorclass）。在这种情况下，func 负责传播其调用到嵌套级别。这允许在将调用传播到嵌套 tensordicts 时进行精细控制。如果为 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)

abstract property names¶

tensordict 的维度名称。

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

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

nanmean(dim: Union[int, Tuple[int], Literal['feature']] = NO_DEFAULT, keepdim: bool = NO_DEFAULT, *, dtype: torch.dtype | None = None) → Any¶

nanmean(dim: Union[int, Tuple[int], Literal['feature']] = NO_DEFAULT, keepdim: bool = NO_DEFAULT, *, dtype: torch.dtype | None = None, reduce: bool) → Union[Any, Tensor]

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

参数:

dim (int, tuple of int, optional) – 如果 None，则返回一个无维度的 tensordict，其中包含所有叶子的平均值（如果可以计算）。如果为整数或整数元组，则仅当此维度与 tensordict 的形状兼容时，才会对指定的维度调用 mean。目前仅允许 “feature” 字符串。使用 dim=”feature” 将实现跨所有特征维度的约简。如果 reduce=True，则将返回一个形状与 TensorDict 的 batch_size 相同的张量。否则，将返回一个具有与 self 相同的结构但特征维度已约简的新 tensordict。
keepdim (bool) – 输出张量是否保留维度。

关键字参数:

dtype (torch.dtype, optional) – 返回张量所需的 dtype。如果指定，则在执行操作之前将输入张量转换为 dtype。这有助于防止数据类型溢出。默认：None。
reduce (bool, optional) – 如果为 True，则会跨所有 TensorDict 值进行归约，并返回一个单一的归约张量。默认为 False。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.nanmean(dim=0)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td.nanmean()
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([]), 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)
>>> td.nanmean(reduce=True)
tensor(-0.0547)
>>> td.nanmean(dim="feature")
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3, 4]),
    device=None,
    is_shared=False)
>>> td = TensorDict(
...     a=torch.ones(3, 4, 5),
...     b=TensorDict(
...         c=torch.ones(3, 4, 5),
...         d=torch.ones(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.nanmean(reduce=True, dim="feature")
tensor([[1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.]])
>>> td.nanmean(reduce=True, dim=0)
tensor([[1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.]])

nansum(dim: Union[int, Tuple[int], Literal['feature']] = NO_DEFAULT, keepdim: bool = NO_DEFAULT, *, dtype: torch.dtype | None = None) → Any¶

nansum(dim: Union[int, Tuple[int], Literal['feature']] = NO_DEFAULT, keepdim: bool = NO_DEFAULT, *, dtype: torch.dtype | None = None, reduce: bool) → Union[Any, Tensor]

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

参数:

dim (int, tuple of int, optional) – 如果 None，则返回一个无维度的 tensordict，其中包含所有叶子的总和（如果可以计算）。如果为整数或整数元组，则仅当此维度与 tensordict 的形状兼容时，才会对指定的维度调用 sum。目前仅允许 “feature” 字符串。使用 dim=”feature” 将实现跨所有特征维度的约简。如果 reduce=True，则将返回一个形状与 TensorDict 的 batch_size 相同的张量。否则，将返回一个具有与 self 相同的结构但特征维度已约简的新 tensordict。
keepdim (bool) – 输出张量是否保留维度。

关键字参数:

dtype (torch.dtype, optional) – 返回张量所需的 dtype。如果指定，则在执行操作之前将输入张量转换为 dtype。这有助于防止数据类型溢出。默认：None。
reduce (bool, optional) – 如果为 True，则会跨所有 TensorDict 值进行归约，并返回一个单一的归约张量。默认为 False。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.nansum(dim=0)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td.nansum()
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([]), 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)
>>> td.nansum(reduce=True)
tensor(-0.)
>>> td.nansum(dim="feature")
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3, 4]),
    device=None,
    is_shared=False)
>>> td = TensorDict(
...     a=torch.ones(3, 4, 5),
...     b=TensorDict(
...         c=torch.ones(3, 4, 5),
...         d=torch.ones(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.nansum(reduce=True, dim="feature")
tensor([[15., 15., 15., 15.],
        [15., 15., 15., 15.],
        [15., 15., 15., 15.]])
>>> td.nansum(reduce=True, dim=0)
tensor([[9., 9., 9., 9., 9.],
        [9., 9., 9., 9., 9.],
        [9., 9., 9., 9., 9.],
        [9., 9., 9., 9., 9.]])

属性 ndim: int¶: 请参阅 batch_dims()。

ndimension() → int¶: 请参阅 batch_dims()。

neg() → Any¶: 计算 TensorDict 中每个元素的 neg() 值。

neg_() → Any¶: 原地计算 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, empty_lazy: bool = False)¶

返回一个大小为 size 的 TensorDict，其中包含空的张量。

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

参数:

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

关键字参数:

dtype (torch.dtype, optional) – 所需返回 tensordict 的类型。默认值：如果 None，则 torch.dtype 将保持不变。
device (torch.device, optional) – 所需返回 tensordict 的设备。默认值：如果 None，则 torch.device 将保持不变。
requires_grad (bool, 可选) – 是否应记录返回张量的自动微分操作。默认值：False。
layout (torch.layout, optional) – 返回的 TensorDict 值的所需布局。默认值：torch.strided。
pin_memory (bool, 可选) – 如果设置为 True，则返回的张量将在固定内存中分配。仅适用于 CPU 张量。默认值：False。
empty_lazy (bool, 可选) – 如果为 True，则将从惰性堆栈中清空其内容。当惰性堆栈的内容可能在填充新 tensordict 时发生更改时，这可能很有用。此参数将传播到子 tensordict。默认为 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, empty_lazy: bool = False)¶

返回一个大小为 size 的 TensorDict，其中填充值为 1。

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

参数:

size (sequence of int) – 定义输出张量形状的整数列表、元组或 torch.Size。
fill_value (scalar) – 用于填充输出 Tensor 的数值。

关键字参数:

dtype (torch.dtype, optional) – 所需返回 tensordict 的类型。默认值：如果 None，则 torch.dtype 将保持不变。
device (torch.device, optional) – 所需返回 tensordict 的设备。默认值：如果 None，则 torch.device 将保持不变。
requires_grad (bool, 可选) – 是否应记录返回张量的自动微分操作。默认值：False。
layout (torch.layout, optional) – 返回的 TensorDict 值的所需布局。默认值：torch.strided。
pin_memory (bool, 可选) – 如果设置为 True，则返回的张量将在固定内存中分配。仅适用于 CPU 张量。默认值：False。
empty_lazy (bool, 可选) – 如果为 True，则将从惰性堆栈中清空其内容。当惰性堆栈的内容可能在填充新 tensordict 时发生更改时，这可能很有用。此参数将传播到子 tensordict。默认为 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, empty_lazy: bool = False)¶

返回一个大小为 size 的 TensorDict，其中填充值为 1。

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

参数:

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

关键字参数:

dtype (torch.dtype, optional) – 所需返回 tensordict 的类型。默认值：如果 None，则 torch.dtype 将保持不变。
device (torch.device, optional) – 所需返回 tensordict 的设备。默认值：如果 None，则 torch.device 将保持不变。
requires_grad (bool, 可选) – 是否应记录返回张量的自动微分操作。默认值：False。
layout (torch.layout, optional) – 返回的 TensorDict 值的所需布局。默认值：torch.strided。
pin_memory (bool, 可选) – 如果设置为 True，则返回的张量将在固定内存中分配。仅适用于 CPU 张量。默认值：False。
empty_lazy (bool, 可选) – 如果为 True，则将从惰性堆栈中清空其内容。当惰性堆栈的内容可能在填充新 tensordict 时发生更改时，这可能很有用。此参数将传播到子 tensordict。默认为 False。

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

使用 data 作为张量返回一个新的 TensorDict。

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

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

参数:

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

关键字参数:

dtype (torch.dtype, optional) – 所需返回 tensordict 的类型。默认值：如果 None，则 torch.dtype 将保持不变。
device (torch.device, optional) – 所需返回 tensordict 的设备。默认值：如果 None，则 torch.device 将保持不变。
requires_grad (bool, 可选) – 是否应记录返回张量的自动微分操作。默认值：False。
pin_memory (bool, 可选) – 如果设置为 True，则返回的张量将在固定内存中分配。仅适用于 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, empty_lazy: bool = False)¶

返回一个大小为 size 的 TensorDict，其中填充值为 0。

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

参数:

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

关键字参数:

dtype (torch.dtype, optional) – 所需返回 tensordict 的类型。默认值：如果 None，则 torch.dtype 将保持不变。
device (torch.device, optional) – 所需返回 tensordict 的设备。默认值：如果 None，则 torch.device 将保持不变。
requires_grad (bool, 可选) – 是否应记录返回张量的自动微分操作。默认值：False。
layout (torch.layout, optional) – 返回的 TensorDict 值的所需布局。默认值：torch.strided。
pin_memory (bool, 可选) – 如果设置为 True，则返回的张量将在固定内存中分配。仅适用于 CPU 张量。默认值：False。
empty_lazy (bool, 可选) – 如果为 True，则将从惰性堆栈中清空其内容。当惰性堆栈的内容可能在填充新 tensordict 时发生更改时，这可能很有用。此参数将传播到子 tensordict。默认为 False。

non_tensor_items(include_nested: bool = False)¶: 递归地返回所有非张量叶节点。

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

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

关键字参数:

out (TensorDict, optional) – 输出 tensordict。
dtype (torch.dtype, optional) – 输出的 dtype (torch>=2.4)。

numel() → int¶

批次中的总元素数量。

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

numpy() → numpy.ndarray | dict[str, Any]¶

将 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!'}}

参见: to_struct_array() 转换为结构化数组。

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

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

关键字参数:: count_duplicates (bool) – 是否将重复的张量计为独立的。如果为 False，则仅丢弃完全相同的张量（来自公共基张量的相同视图但不同的 ID 将被计数两次）。默认为 True（每个张量假定为单个副本）。

permute(*dims: int)¶

permute(dims: list | tuple)

根据 dims 返回 tensordict 的视图，其中批次维度已重新排序。

参数:

*dims_list (int) – tensordict 的批处理维度的顺序。或者，可以提供单个整数可迭代对象。
dims (list of int) – 调用 permute(…) 的替代方法。

返回:

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

示例

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

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

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

参数:

num_threads (int 或 str) – 如果提供，则用于调用 pin_memory 的线程数。默认为 None，它在 ThreadPoolExecutor(max_workers=None) 中设置了高线程数。要对所有 pin_memory() 的调用在主线程上执行，请传入 num_threads=0。
inplace (bool, optional) – 如果为 True，则 tensordict 将就地修改。默认为 False。

pin_memory_(num_threads: int | str = 0) → Any¶

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

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

pop(key: NestedKey, default: Any = _NoDefault.ZERO) → Union[Tensor, TensorCollection]¶

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

如果未找到该值且未提供默认值，则会引发 KeyError。

参数:

key (str or nested key) – 要查找的条目。
default (Any, optional) – 找不到键时要返回的值。

示例

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

抽象 popitem() → Tuple[NestedKey, Union[Tensor, TensorCollection]]¶

移除最后插入 TensorDict 的项。

popitem 只会返回非嵌套值。

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

将 self 的每个元素与 other 求幂，并返回结果张量。

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

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

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

pow_(other: tensordict.base.TensorDictBase | torch.Tensor) → Any¶: 原地版本的 pow()。

注意

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

prod(dim: Union[int, Tuple[int], Literal['feature']] = NO_DEFAULT, keepdim: bool = NO_DEFAULT, *, dtype: torch.dtype | None = None) → Any¶

prod(dim: Union[int, Tuple[int], Literal['feature']] = NO_DEFAULT, keepdim: bool = NO_DEFAULT, *, dtype: torch.dtype | None = None, reduce: bool) → Union[Any, Tensor]

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

参数:

dim (int, tuple of int, optional) – 如果 None，则返回一个无维度的 tensordict，其中包含所有叶子的积（如果可以计算）。如果为整数或整数元组，则仅当此维度与 tensordict 的形状兼容时，才会对指定的维度调用 prod。目前仅允许 “feature” 字符串。使用 dim=”feature” 将实现跨所有特征维度的约简。如果 reduce=True，则将返回一个形状与 TensorDict 的 batch_size 相同的张量。否则，将返回一个具有与 self 相同的结构但特征维度已约简的新 tensordict。
keepdim (bool) – 输出张量是否保留维度。

关键字参数:

dtype (torch.dtype, optional) – 返回张量所需的 dtype。如果指定，则在执行操作之前将输入张量转换为 dtype。这有助于防止数据类型溢出。默认：None。
reduce (bool, optional) – 如果为 True，则会跨所有 TensorDict 值进行归约，并返回一个单一的归约张量。默认为 False。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.prod(dim=0)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td.prod()
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([]), 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)
>>> td.prod(reduce=True)
tensor(-0.)
>>> td.prod(dim="feature")
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3, 4]),
    device=None,
    is_shared=False)
>>> td = TensorDict(
...     a=torch.ones(3, 4, 5),
...     b=TensorDict(
...         c=torch.ones(3, 4, 5),
...         d=torch.ones(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.prod(reduce=True, dim="feature")
tensor([[1., 1., 1., 1.],
        [1., 1., 1., 1.],
        [1., 1., 1., 1.]])
>>> td.prod(reduce=True, dim=0)
tensor([[1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.],
        [1., 1., 1., 1., 1.]])

qint32() → Any¶: 将所有张量强制转换为 torch.qint32。

qint8() → Any¶: 将所有张量强制转换为 torch.qint8。

quint4x2() → Any¶: 将所有张量强制转换为 torch.quint4x2。

quint8() → Any¶: 将所有张量强制转换为 torch.quint8。

reciprocal() → Any¶: 计算 TensorDict 中每个元素的 reciprocal() 值。

reciprocal_() → Any¶: 原地计算 TensorDict 中每个元素的 reciprocal() 值。

record_stream(stream: Stream) → Any¶

标记 TensorDict 已被此流使用。

当 TensorDict 被解除分配时，请确保在解除分配时该流上的所有工作完成之前，不要将张量内存重用于其他张量。

更多信息请参阅 record_stream()。

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

接收 TensorDict 的内容并用其更新内容。

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

参数:

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

关键字参数:

group (torch.distributed.ProcessGroup, 可选) – 如果设置，将使用指定的进程组进行通信。否则，将使用默认进程组。默认为 `None`。
init_tag (int) – 用于标记张量的 init_tag。请注意，此值将按 TensorDict 中包含的张量数量进行递增。
pseudo_rand (bool) – 如果为 True，则标签序列将是伪随机的，允许不同节点之间的数据发送而不会重叠。请注意，这些伪随机数的生成成本很高（每秒 1e-5），这意味着它可能会减慢算法的运行速度。此值必须与传递给 send() 的值匹配。默认为 False。

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

在所有机器上归约 TensorDict。

只有具有 rank dst 的进程才能接收最终结果。

refine_names(*names) → Any¶

根据名称精炼 self 的维度名称。

精炼是重命名的特殊情况，它“提升”未命名维度。None 维度可以精炼为任何名称；已命名维度只能精炼为相同名称。

由于命名张量可以与未命名张量共存，因此细化名称提供了一种很好的方式来编写既适用于命名张量又适用于未命名张量的命名张量感知代码。

names 可以包含最多一个 Ellipsis (...)。Ellipsis 会贪婪地展开；它会就地展开以使 names 的长度与 self.dim() 相同，方法是使用 self.names 的相应索引处的名称。

Returns: 具有根据输入命名的维度的相同的 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) → Any¶

用新字符串重命名一个键，并返回具有更新后的键名的相同 TensorDict。

参数:

old_key (str 或 nested key) – 要重命名的键。
new_key (str 或 nested key) – 条目的新名称。
safe (bool, optional) – 如果为 True，则当新键已存在于 TensorDict 中时将引发错误。

返回:

self

repeat(repeats: Size)¶

沿指定维度重复此张量。

与 expand() 不同，此函数会复制张量的数据。

警告

repeat() 的行为与 repeat() 不同，但更接近于 numpy.tile()。对于与 numpy.repeat() 类似的运算符，请参见 repeat_interleave()。

参数:: repeat (torch.Size, int..., tuple of int or list of int) – 沿每个维度重复此张量的次数。

示例

>>> import torch
>>>
>>> from tensordict import TensorDict
>>>
>>> td = TensorDict(
...     {
...         "a": torch.randn(3, 4, 5),
...         "b": TensorDict({
...             "c": torch.randn(3, 4, 10, 1),
...             "a string": "a string!",
...         }, batch_size=[3, 4, 10])
...     }, batch_size=[3, 4],
... )
>>> print(td.repeat(1, 2))
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 8, 5]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                a string: NonTensorData(data=a string!, batch_size=torch.Size([3, 8, 10]), device=None),
                c: Tensor(shape=torch.Size([3, 8, 10, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 8, 10]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3, 8]),
    device=None,
    is_shared=False)

抽象 repeat_interleave(repeats: torch.Tensor | int, dim: Optional[int] = None, *, output_size: Optional[int] = None) → Any¶

重复 TensorDict 的元素。

警告

这与 repeat() 不同，但与 numpy.repeat() 类似。

参数:

repeats (torch.Tensor or int) – 每个元素的重复次数。repeats 将被广播以适应给定轴的形状。
dim (int, optional) – 重复值的维度。默认为使用展平的输入数组，并返回一个展平的输出数组。

关键字参数:

output_size (int, 可选) – 指定轴的总输出大小（例如，重复次数的总和）。如果给定，它将避免同步流以计算 tensordict 的输出形状。

返回:

重复的 TensorDict，其形状与输入相同，除了在给定轴上。

示例

>>> import torch
>>>
>>> from tensordict import TensorDict
>>>
>>> td = TensorDict(
...     {
...         "a": torch.randn(3, 4, 5),
...         "b": TensorDict({
...             "c": torch.randn(3, 4, 10, 1),
...             "a string": "a string!",
...         }, batch_size=[3, 4, 10])
...     }, batch_size=[3, 4],
... )
>>> print(td.repeat_interleave(2, dim=0))
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([6, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                a string: NonTensorData(data=a string!, batch_size=torch.Size([6, 4, 10]), device=None),
                c: Tensor(shape=torch.Size([6, 4, 10, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([6, 4, 10]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([6, 4]),
    device=None,
    is_shared=False)

replace(*args, **kwargs)¶

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

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

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

requires_grad_(requires_grad=True) → Any¶

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

返回此 TensorDict。

参数:: requires_grad (bool, optional) – autograd 是否应该记录此 tensordict 上的操作。默认为 True。

抽象 reshape(*shape: int)¶

抽象 reshape(shape: list | tuple)

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

参数:: *shape (int) – 结果 TensorDict 的新形状。
返回:: 具有重塑键的 TensorDict。

示例

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

round() → Any¶: 计算 TensorDict 中每个元素的 round() 值。

round_() → Any¶: 原地计算 TensorDict 中每个元素的 round() 值。

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

从 self 中减去 other 乘以 alpha 的缩放值。

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

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

参数:

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

关键字参数:

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

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

将tensordict保存到磁盘。

此函数是 memmap() 的代理。

属性 saved_path¶

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

此参数在 is_memmap() 返回 False（例如，当 TensorDict 解锁时）后即失效。

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

选择 TensorDict 的键，并返回一个仅包含所选键的新 TensorDict。

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

参数:

*keys (str) – 要选择的键。
inplace (bool) – 如果为 True，则 tensordict 将就地剪枝。默认为 False。
strict (bool, optional) – 选择一个不存在的键是否会返回错误。默认为 True。

返回:

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

注意

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

示例

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

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

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

参数:

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

关键字参数:

group (torch.distributed.ProcessGroup, 可选) – 如果设置，将使用指定的进程组进行通信。否则，将使用默认进程组。默认为 `None`。
init_tag (int) – 用于标记张量的初始标签。注意，该标签将按 TensorDict 中包含的张量数量进行递增。
pseudo_rand (bool) – 如果为 True，则标签序列将是伪随机的，允许来自不同节点的多数据发送而不发生重叠。请注意，这些伪随机数的生成成本很高（每秒 1e-5），这意味着它可能会减慢算法的运行速度。默认为 False。

示例

>>> from torch import multiprocessing as mp
>>> from tensordict import TensorDict
>>> import torch
>>>
>>>
>>> def client():
...     torch.distributed.init_process_group(
...         "gloo",
...         rank=1,
...         world_size=2,
...         init_method=f"tcp://: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://: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()

separates(*keys: NestedKey, default: Any = _NoDefault.ZERO, strict: bool = True, filter_empty: bool = True) → Any¶

将指定的键从 tensordict 中分离（就地操作）。

另请参阅

此方法等同于在单个分割上使用 inplace=True 调用 split_keys()。

另请参阅

此方法等同于 exclude()，但它返回数据的另一个分割。

参数:

keys (NestedKey) – 要从 tensordict 中分离的键。
default (Any, optional) – 当键丢失时返回的值。如果未指定且 strict=True，则会引发异常。否则，任何缺失键的默认值将是 None，除非另有指定。
strict (bool, optional) – 如果为 True，则在键丢失时引发异常。默认为 True。
filter_empty (bool, optional) – 如果为 True，则会删除 self 中的空 tensordicts。默认为 True。

返回:

分离出的 tensordict。

返回类型:

T

示例

>>> td = TensorDict(
...     a=0,
...     b=0,
...     c=0,
...     d=0,
... )
>>> td_a_c = td.separates("a", "c")
>>> print(td_a_c)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False),
        c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.int64, is_shared=False)},
    batch_size=torch.Size([]),
    device=None,
    is_shared=False)
>>> print(td)
TensorDict(
    fields={
        b: 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)

set(key: NestedKey, item: Union[Tensor, TensorCollection], inplace: bool = False, *, non_blocking: bool = False, **kwargs: Any) → Any¶

设置一个新的键值对。

参数:

key (str, str 元组) – 要设置的键的名称。
item (torch.Tensor or equivalent, TensorDictBase instance) – 要存储在 tensordict 中的值。
inplace (bool, 可选) – 如果为 True 且键匹配 tensordict 中的现有键，则更新将在该键值对上原地发生。如果 inplace 为 True 且找不到条目，则会添加该条目。要进行更严格的原地操作，请改用 set_()。默认为 False。

关键字参数:

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

返回:

self

示例

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

set_(key: NestedKey, item: Union[Tensor, TensorCollection], *, non_blocking: bool = False) → Any¶

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

参数:

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

关键字参数:

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

返回:

self

示例

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

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

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

参数:

key (str, str 元组) – 要修改的键。
value (torch.Tensor) – 在 index 处设置的值。
index (int, tensor or tuple) – 写入值的索引。

关键字参数:

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

返回:

self

示例

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

set_non_tensor(key: NestedKey, value: Any)¶

使用 tensordict.tensorclass.NonTensorData 在 tensordict 中注册一个非张量值。

可以使用 TensorDictBase.get_non_tensor() 或直接使用 get 检索该值，后者将返回 tensordict.tensorclass.NonTensorData 对象。

返回：self

示例

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

setdefault(key: NestedKey, default: Union[Tensor, TensorCollection, Any], inplace: bool = False) → Union[Tensor, TensorCollection]¶

如果 key 不在 tensordict 中，则将 key 条目插入其值为 default。

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

参数:

key (str 或 nested key) – 值的名称。
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。

abstract share_memory_() → Any¶

将所有张量放入共享内存。

TensorDict 将被锁定，这意味着任何非原地写入操作都会引发异常（例如，重命名、设置或删除条目）。反之，一旦 tensordict 被解锁，share_memory 属性将变为 False，因为进程间的身份不再有保证。

返回:: self

sigmoid() → Any¶: 计算 TensorDict 中每个元素的 sigmoid() 值。

sigmoid_() → Any¶: 原地计算 TensorDict 中每个元素的 sigmoid() 值。

sign() → Any¶: 计算 TensorDict 中每个元素的 sign() 值。

sign_() → Any¶: 原地计算 TensorDict 中每个元素的 sign() 值。

sin() → Any¶: 计算 TensorDict 中每个元素的 sin() 值。

sin_() → Any¶: 原地计算 TensorDict 中每个元素的 sin() 值。

sinh() → Any¶: 计算 TensorDict 中每个元素的 sinh() 值。

sinh_() → Any¶: 原地计算 TensorDict 中每个元素的 sinh() 值。

size(dim: Optional[int] = None) → torch.Size | int¶

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

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

softmax(dim: int, dtype: Optional[dtype] = None)¶

将 softmax 函数应用于 tensordict 元素。

参数:

dim (int 或 int 元组) – 将在其上计算 softmax 的 tensordict 维度。
dtype (torch.dtype, optional) – 返回张量的期望数据类型。如果指定，则在执行操作之前，输入张量将被转换为 dtype。这有助于防止数据类型溢出。

property sorted_keys: list[tensordict._nestedkey.NestedKey]¶

按字母顺序返回键。

不支持其他参数。

如果 TensorDict 被锁定，键将被缓存，直到 TensorDict 解锁以获得更快的执行。

abstract split(split_size: int | list[int], dim: int = 0) → list[tensordict._tensorcollection.TensorCollection]¶

像 torch.split 一样，使用指定的尺寸在给定维度上分割 TensorDict 中的每个张量。

返回一个包含分割块视图的 TensorDict 实例列表。

参数:

split_size (int or 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, default: Any = _NoDefault.ZERO, strict: bool = True, reproduce_struct: bool = False) → Tuple[T, ...]¶

根据一个或多个键集将 tensordict 分割成子集。

该方法将返回 N+1 个 tensordict，其中 N 是提供的参数数量。

参数:

key_sets (键的序列 Dict[in_key, out_key] 或 键列表) – 各个分割。
inplace (bool, optional) – 如果为 True，则 self 中的键将被原地移除。默认为 False。
default (Any, 可选) – 当 key 丢失时要返回的值。如果未指定且 strict=True，则会引发异常。
strict (bool, optional) – 如果为 True，则在键丢失时引发异常。默认为 True。
reproduce_struct (bool, optional) – 如果为 True，则所有返回的 tensordict 都将具有与 self 相同的树结构，即使某些子 tensordict 不包含叶子。

注意

None 非张量值将被忽略且不返回。

注意

该方法不检查提供的列表中的重复项。

示例

>>> td = TensorDict(
...     a=0,
...     b=0,
...     c=0,
...     d=0,
... )
>>> td_a, td_bc, td_d = td.split_keys(["a"], ["b", "c"])
>>> print(td_bc)

sqrt() → Any¶: 计算 self 的逐元素平方根。

sqrt_() → Any¶: sqrt() 的原地版本。

squeeze(dim: int | None = None) → Self¶

压缩维度在 -self.batch_dims+1 和 self.batch_dims-1 之间的所有张量，并将它们返回到一个新的 tensordict 中。

参数:: dim (Optional[int]) – 压缩的维度。如果 dim 为 None，则会压缩所有单例维度。默认为 None。

示例

>>> td = TensorDict({
...     'x': torch.arange(24).reshape(3, 1, 4, 2),
... }, batch_size=[3, 1, 4])
>>> td = td.squeeze()
>>> td.shape
torch.Size([3, 4])
>>> td.get("x").shape
torch.Size([3, 4, 2])

此操作也可以作为上下文管理器使用。对原始 tensordict 的更改将是异变的，即原始张量的内容不会被修改。这也假定 tensordict 未被锁定（否则，需要解锁 tensordict）。此功能与隐式压缩不兼容。

>>> td = TensorDict({
...     'x': torch.arange(24).reshape(3, 1, 4, 2),
... }, batch_size=[3, 1, 4])
>>> with td.squeeze(1) as tds:
...     tds.set("y", torch.zeros(3, 4))
>>> assert td.get("y").shape == [3, 1, 4]

classmethod stack(input, dim: int = 0, *, out=None)¶

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

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

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

将 tensordict 的所有条目堆叠成一个张量。

参数:

dim (int, optional) – 条目应堆叠的维度。

关键字参数:

sorted (bool 或 NestedKeys 列表) – 如果为 True，则条目将按字母顺序堆叠。如果为 False（默认值），则使用字典顺序。或者，可以提供键名列表，并相应地堆叠张量。这会产生一些开销，因为键列表将与 tensordict 中的叶名称列表进行比较。
out (torch.Tensor, optional) – 用于堆叠操作的可选目标张量。

stack_tensors(*keys: NestedKey, out_key: NestedKey, dim: int = 0, keep_entries: bool = False) → Any¶

将条目堆叠到新的条目中，并可能移除原始值。

参数:

keys (NestedKey 序列) – 要堆叠的条目。

关键字参数:

out_key (NestedKey) – 堆叠输入的新的键名。
keep_entries (bool, 可选) – 如果为 False，则 keys 中的条目将被删除。默认为 False。
dim (int, 可选) – 堆叠发生的维度。默认为 0。

返回: self

示例

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

state_dict(destination=None, prefix='', keep_vars=False, flatten=False) → OrderedDict[str, Any]¶

从 tensordict 中生成 state_dict。

state_dict 的结构将保持嵌套，除非将 flatten 设置为 True。

tensordict state_dict 包含重建 tensordict 所需的所有张量和元数据（目前不支持名称）。

参数:

destination (dict, optional) – 如果提供，tensordict 的状态将更新到 dict 中，并返回相同的对象。否则，将创建一个 OrderedDict 并返回。默认：None。
prefix (str, optional) – 添加到张量名称的前缀，用于组合 state_dict 中的键。默认：''。
keep_vars (bool, optional) – 默认情况下，state_dict 中返回的 torch.Tensor 项将从 autograd 中分离。如果设置为 True，则不会执行分离。默认：False。
flatten (bool, optional) – 是否应使用 "." 字符展平结构。默认为 False。

示例

>>> data = TensorDict({"1": 1, "2": 2, "3": {"3": 3}}, [])
>>> sd = data.state_dict()
>>> print(sd)
OrderedDict([('1', tensor(1)), ('2', tensor(2)), ('3', OrderedDict([('3', tensor(3)), ('__batch_size', torch.Size([])), ('__device', None)])), ('__batch_size', torch.Size([])), ('__device', None)])
>>> sd = data.state_dict(flatten=True)
OrderedDict([('1', tensor(1)), ('2', tensor(2)), ('3.3', tensor(3)), ('__batch_size', torch.Size([])), ('__device', None)])

std(dim: Union[int, Tuple[int], Literal['feature']] = NO_DEFAULT, keepdim: bool = NO_DEFAULT, *, correction: int = 1) → Any¶

std(dim: Union[int, Tuple[int], Literal['feature']] = NO_DEFAULT, keepdim: bool = NO_DEFAULT, *, correction: int = 1, reduce: bool) → Union[Any, Tensor]

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

参数:

dim (int, tuple of int, optional) – 如果为 None，则返回一个无量纲的 tensordict，其中包含所有叶子节点的总和值（如果可计算）。如果为整数或整数元组，则仅当该维度与 tensordict 形状兼容时，才会在指定的维度上调用 std。目前只允许 “feature” 字符串。使用 dim=”feature” 将在所有特征维度上进行归约。如果 reduce=True，将返回一个形状与 TensorDict 的 batch-size 相同的张量。否则，将返回一个结构与 self 相同的、具有归约特征维度的新 tensordict。
keepdim (bool) – 输出张量是否保留维度。

关键字参数:

correction (int) – 样本大小与样本自由度之间的差值。默认为贝塞尔校正，correction=1。
reduce (bool, optional) – 如果为 True，则会跨所有 TensorDict 值进行归约，并返回一个单一的归约张量。默认为 False。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.std(dim=0)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td.std()
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([]), 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)
>>> td.std(reduce=True)
tensor(1.0006)
>>> td.std(dim="feature")
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3, 4]),
    device=None,
    is_shared=False)
>>> td = TensorDict(
...     a=torch.ones(3, 4, 5),
...     b=TensorDict(
...         c=torch.ones(3, 4, 5),
...         d=torch.ones(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.std(reduce=True, dim="feature")
tensor([[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]])
>>> td.std(reduce=True, dim=0)
tensor([[0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.]])

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

将 other 缩放 alpha 后，从 self 中减去。

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

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

参数:

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

关键字参数:

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

sub_(other: tensordict.base.TensorDictBase | torch.Tensor | float, alpha: float | None = None)¶: sub() 的原地版本。

注意

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

sum(dim: Union[int, Tuple[int], Literal['feature']] = NO_DEFAULT, keepdim: bool = NO_DEFAULT, *, dtype: torch.dtype | None = None) → Any¶

sum(dim: Union[int, Tuple[int], Literal['feature']] = NO_DEFAULT, keepdim: bool = NO_DEFAULT, *, dtype: torch.dtype | None = None, reduce: bool) → Union[Any, Tensor]

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

参数:

dim (int, tuple of int, optional) – 如果 None，则返回一个无维度的 tensordict，其中包含所有叶子的总和（如果可以计算）。如果为整数或整数元组，则仅当此维度与 tensordict 的形状兼容时，才会对指定的维度调用 sum。目前仅允许 “feature” 字符串。使用 dim=”feature” 将实现跨所有特征维度的约简。如果 reduce=True，则将返回一个形状与 TensorDict 的 batch_size 相同的张量。否则，将返回一个具有与 self 相同的结构但特征维度已约简的新 tensordict。
keepdim (bool) – 输出张量是否保留维度。

关键字参数:

dtype (torch.dtype, optional) – 返回张量所需的 dtype。如果指定，则在执行操作之前将输入张量转换为 dtype。这有助于防止数据类型溢出。默认：None。
reduce (bool, optional) – 如果为 True，则会跨所有 TensorDict 值进行归约，并返回一个单一的归约张量。默认为 False。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.sum(dim=0)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td.sum()
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([]), 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)
>>> td.sum(reduce=True)
tensor(-0.)
>>> td.sum(dim="feature")
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3, 4]),
    device=None,
    is_shared=False)
>>> td = TensorDict(
...     a=torch.ones(3, 4, 5),
...     b=TensorDict(
...         c=torch.ones(3, 4, 5),
...         d=torch.ones(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.sum(reduce=True, dim="feature")
tensor([[15., 15., 15., 15.],
        [15., 15., 15., 15.],
        [15., 15., 15., 15.]])
>>> td.sum(reduce=True, dim=0)
tensor([[9., 9., 9., 9., 9.],
        [9., 9., 9., 9., 9.],
        [9., 9., 9., 9., 9.],
        [9., 9., 9., 9., 9.]])

tan() → Any¶: 计算 TensorDict 中每个元素的 tan() 值。

tan_() → Any¶: 原地计算 TensorDict 中每个元素的 tan() 值。

tanh() → Any¶: 计算 TensorDict 中每个元素的 tanh() 值。

tanh_() → Any¶: 原地计算 TensorDict 中每个元素的 tanh() 值。

tensor_split(indices_or_sections: int | list[int] | tuple[int, ...] | torch.Tensor, dim=0) → tuple[tensordict.base.TensorDictBase, ...]¶

将 TensorDict 沿着 dim 维度根据 indices_or_sections 指定的索引或节数分割成多个子 tensordict，所有这些子 tensordict 都是输入张量的视图。

参数:

indices_or_sections (int 或 List(int) 或 tuple(int) 或 1D tensor of ints) – 如果 indices_or_sections 是一个整数 n 或值为 n 的零维长整型张量，则输入沿 dim 维度被分成 n 部分。如果输入在 dim 维度上可被 n 整除，则每部分的长度相等，为 input.size(dim) / n。如果输入不能被 n 整除，则前 int(input.size(dim) % n) 部分的长度将为 int(input.size(dim) / n) + 1，其余部分的长度为 int(input.size(dim) / n)。如果 indices_or_sections 是一个整数列表或元组，或是一维长整型张量，则输入将在 dim 维度上，在列表、元组或张量中的每个索引处进行分割。例如，indices_or_sections=[2, 3] 和 dim=0 将产生张量 input[:2]、input[2:3] 和 input[3:]。如果 indices_or_sections 是一个张量，它必须是一个在 CPU 上的零维或一维长整型张量。
dim (int, optional) – 要分割张量的维度。默认为 0

示例

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

to(device: Optional[Union[int, device]] = ..., dtype: Optional[Union[device, str]] = ..., non_blocking: bool = ..., inplace: bool = False) → Any¶

to(dtype: Union[device, str], non_blocking: bool = ...) → Any

to(tensor: Tensor, non_blocking: bool = ...) → Any

to(*, other: T, non_blocking: bool = ...) → Any

to(*, batch_size: Size) → Any

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

不允许将张量转换为新的 dtype，因为 tensordict 不会绑定到只包含单一 tensor dtype。

参数:

device (torch.device, optional) – tensordict 的期望设备。
dtype (torch.dtype, optional) – tensordict 的期望浮点数或复数 dtype。
tensor (torch.Tensor, optional) – 此 TensorDict 中所有张量的期望 dtype 和设备都取自此张量。

关键字参数:

non_blocking (bool, 可选) – 操作是否应该阻塞。
memory_format (torch.memory_format, optional) – 此 tensordict 中 4D 参数和缓冲区所需的内存格式。
batch_size (torch.Size, optional) – 输出 tensordict 的结果 batch-size。
other (TensorDictBase, 可选) –
TensorDict 实例，其 dtype 和 device 是该 TensorDict 中所有张量的目标 dtype 和 device。

注意

由于 TensorDictBase 实例没有 dtype，因此 dtype 是从示例叶节点收集的。如果有多个 dtype，则不进行类型转换。
non_blocking_pin (bool, 可选) –
如果为 True，则张量将在发送到设备之前被 pinned。这将异步进行，但可以通过 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() 调用中的任何多线程。
inplace (bool, 可选) – 如果为 True，则数据将被就地写入同一个 tensordict。每当构建 tensordict 受 CPU 开销限制时，这可以显著加快速度。默认为 False。

返回:

如果设备与 tensordict 设备不同，以及/或传递了 dtype，则返回一个新的 tensordict 实例。否则，返回相同的 tensordict。batch_size 的修改是就地进行的。

注意

如果 TensorDict 是合并的，结果 TensorDict 也将是合并的。每个新张量都是映射到目标设备的合并存储的视图。

示例

>>> data = TensorDict({"a": 1.0}, [], device=None)
>>> data_cuda = data.to("cuda:0")  # casts to cuda
>>> data_int = data.to(torch.int)  # casts to int
>>> data_cuda_int = data.to("cuda:0", torch.int)  # multiple casting
>>> data_cuda = data.to(torch.randn(3, device="cuda:0"))  # using an example tensor
>>> data_cuda = data.to(other=TensorDict({}, [], device="cuda:0"))  # using a tensordict example

to_dict(*, retain_none: bool = True, convert_tensors: Union[bool, Literal['numpy']] = False, tolist_first: bool = False) → dict[str, Any]¶

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

参数:

retain_none (bool) – 如果为 True，则 tensorclass 实例中的 None 值将被写入字典。否则，它们将被丢弃。默认值：True。
convert_tensors (bool, "numpy") – 如果为 True，则在创建字典时，张量将被转换为列表。如果为“numpy”，则张量将被转换为 numpy 数组。否则，它们将保持为张量。默认值：False。
tolist_first (bool) – 如果为 True，则当 tensordict 具有批处理维度时，它将首先转换为列表。默认值：False。

返回:

tensordict 的字典表示。

另请参阅

tolist()

示例

>>> import torch
>>> from tensordict import TensorDict
>>>
>>> td = TensorDict(
...     a=torch.arange(24).view(2, 3, 4),
...     b=TensorDict(c=torch.arange(12).reshape(2, 3, 2), batch_size=(2, 3, 2)),
...     batch_size=(2, 3)
... )
>>> print(td.to_dict())
{'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]]]), ‘b’: {‘c’: tensor([[[ 0, 1], [ 2, 3], [ 4, 5]],

[[ 6, 7],
[ 8, 9], [10, 11]]])}}

>>> print(td.to_dict(convert_tensors=True))
{'a': [[[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11]], [[12, 13, 14, 15], [16, 17, 18, 19], [20, 21, 22, 23]]], 'b': {'c': [[[0, 1], [2, 3], [4, 5]], [[6, 7], [8, 9], [10, 11]]]}}

to_h5(filename, **kwargs) → Any¶

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

参数:

filename (str 或 path) – h5 文件的路径。
**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_lazystack(dim: int = 0)¶

将 TensorDict 转换为 LazyStackedTensorDict 或等价物。

注意

此方法可用于更改 LazyStackedTensorDict 的堆叠维度。例如，如果您有一个 stack_dim=1 的 LazyStackedTensorDict，您可以使用此方法将其更改为 stack_dim=0。

>>> td = TensorDict({"a": torch.zeros(2, 3), "b": torch.ones(2, 3)}, batch_size=(2, 3))
>>> td2 = td.to_lazystack()
>>> td2.batch_size
torch.Size([2, 3])
>>> assert isinstance(td2, LazyStackedTensorDict)
>>> assert td2.stack_dim == 0
>>> td3 = td2.to_lazystack(1)
>>> assert td3.stack_dim == 1
>>> td3.batch_size
torch.Size([2, 3])

参数:: dim (int, optional) – 要堆叠 tensordict 的维度。默认为 0。
返回:: LazyStackedTensorDict 实例。

示例

>>> from tensordict import TensorDict
>>> td = TensorDict({"a": torch.zeros(2, 3), "b": torch.ones(2, 3)}, batch_size=(2, 3))
>>> td2 = td.to_lazystack()
>>> td2.batch_size
torch.Size([2, 3])
>>> assert isinstance(td2, LazyStackedTensorDict)

将 TensorCollection 的内容写入流式数据集。

关键字参数:

out (str | Tuple[str, str]) –
用于保存分片文件的输出数据集目录。
1. 如果 out 是本地目录，则分片文件将本地保存。
2. 如果 out 是远程目录，则将创建一个本地临时目录来
  缓存分片文件，然后将分片文件上传到远程位置。完成后，在分片上传后删除临时目录。
3. 如果 out 是 (local_dir, remote_dir) 的元组，则分片文件将保存在
  local_dir 中，并上传到远程位置。
columns (dict[str, str]) – 可选的列字典。如果未提供，将自动从张量集合中推断。
keep_local (bool) – 如果数据集已上传，则在上传后是否保留本地数据集目录。默认为 False。
compression (str, optional) – 可选的压缩或压缩级别。默认为 None。
hashes (List[str], optional) – 可选的要应用于分片文件的哈希算法列表。默认为 None。
size_limit (Union[int, str], 可选) – 可选的分片大小限制，在此之后开始新的分片。如果为 None，则将所有内容放入一个分片。也可以指定字节人类可读格式，例如 "100kb" 表示 100 千字节 (100*1024) 等。默认为 1 << 26。如果
writer – (MDSWriter, optional): 要使用的写入器。将从 out 关键字参数和其他输入关键字参数创建。
**kwargs (Any) – 写入器的其他设置。

注意

MDSWriter 对嵌套字典的支持有限。处理嵌套 tensordict 的正确方法是先使用 flatten_keys() 方法进行写入，然后在读取后使用 unflatten_keys() 方法。

警告

此方法需要安装 mosaicml-streaming。

警告

对于非张量数据，数据类型必须是固定的。tensordict 恢复数据类型的方式是查看列表的第一个元素。如果为 None，将引发错误。否则，列表中的所有数据都必须具有相同的类型（或为 None 表示缺失）。

另请参阅

请参阅 Mosaic streaming 库 API https://docs.mosaicml.com/projects/streaming

以下示例展示了如何创建数据集并在 PyTorch dataloader 中加载它的端到端示例。

示例

>>> import tempfile
>>> from typing import Any
>>> from tensordict import TensorDict, LazyStackedTensorDict
>>> import torch
>>>
>>>
>>> td = LazyStackedTensorDict(
...     TensorDict(a=0, b=1, c=torch.randn(2), d="a string"),
...     TensorDict(a=0, b=1, c=torch.randn(3), d="another string"),
...     TensorDict(a=0, b=1, c=torch.randn(3), d="yet another string"),
... )
>>>
>>> with tempfile.TemporaryDirectory() as tempdir:
...     # Create a dataset on one process / thread / node...
...     td.to_mds(out=tempdir)
...
...     # Create a dataloader
...     from streaming import StreamingDataset
...
...     # Load the dataset on another thread / process / node...
...     dataset = StreamingDataset(local=tempdir, remote=None, batch_size=2)
...
...     # Use the class `from_list` method as a collate_fn
...     dl = torch.utils.data.DataLoader(dataset=dataset, batch_size=2, collate_fn=LazyStackedTensorDict.from_list)
...     for batch in dl:
...         print("batch", batch)

to_module(module: Module, *, inplace: bool | None = None, return_swap: bool = True, swap_dest=None, use_state_dict: bool = False, non_blocking: bool = False, memo=None)¶

将 TensorDictBase 实例的内容递归地写入指定的 nn.Module 属性。

to_module 也可以用作上下文管理器，用于临时用一组参数/缓冲区填充模块（请参阅下面的示例）。

参数:

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

关键字参数:

inplace (bool, 可选) – 如果为 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.data.zero_()
>>> params.to_module(module)
>>> assert (module.layers[0].linear1.weight == 0).all()

将 tensordict 用作上下文管理器可以方便地进行函数调用： .. rubric:: 示例

>>> from tensordict import from_module
>>> module = nn.TransformerDecoder(
...     decoder_layer=nn.TransformerDecoderLayer(nhead=4, d_model=4),
...     num_layers=1)
>>> params = TensorDict.from_module(module)
>>> params = params.data * 0 # Use TensorDictParams to remake these tensors regular nn.Parameter instances
>>> with params.to_module(module):
...     # Call the module with zeroed params
...     y = module(*inputs)
>>> # The module is repopulated with its original params
>>> assert (TensorDict.from_module(module) != 0).any()

返回:: 一个 tensordict，如果 return_swap 为 True，则包含来自模块的值，否则为 None。

to_namedtuple(dest_cls: Optional[type] = None) → Any¶

将 tensordict 转换为命名元组。

参数:: dest_cls (Type, optional) – 一个可选的命名元组类。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> data = TensorDict({
...     "a_tensor": torch.zeros((3)),
...     "nested": {"a_tensor": torch.zeros((3)), "a_string": "zero!"}}, [3])
>>> data.to_namedtuple()
GenericDict(a_tensor=tensor([0., 0., 0.]), nested=GenericDict(a_tensor=tensor([0., 0., 0.]), a_string='zero!'))

to_padded_tensor(padding=0.0, mask_key: Optional[NestedKey] = None) → Any¶

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

参数:

padding (float) – tensordict 中张量的填充值。默认为 0.0。
mask_key (NestedKey, optional) – 如果提供，则为写入有效值掩码的键。如果异构维度不是 tensordict batch-size 的一部分，则会导致错误。默认为 None

to_pytree()¶

将 tensordict 转换为 PyTree。

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

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

to_struct_array() → ndarray¶

将 tensordict 转换为 numpy 结构化数组。

在 from_struct_array() - to_struct_array() 循环中，输入和输出数组的内容应匹配。但是，to_struct_array 不会保留原始数组的内存内容。

另请参阅

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

另请参阅

有关转换为 numpy 数组字典，请参阅 numpy()。

返回:: 输入 TensorDict 的 numpy 结构化数组表示。

示例

>>> import torch
>>> from tensordict import TensorDict
>>> td = TensorDict({'a': torch.tensor([1, 2, 3]), 'b': torch.tensor([4.0, 5.0, 6.0])}, batch_size=[3])
>>> arr = td.to_struct_array()
>>> print(arr)
[(1, 4.) (2, 5.) (3, 6.)]

to_tensordict(*, retain_none: Optional[bool] = None) → Any¶

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

参数:: retain_none (bool) – 如果为 True，则 tensorclass 实例中的 None 值将被写入 tensordict。否则，它们将被丢弃。默认值：True。
返回:: 包含相同值的新 TensorDict 对象。

tolist(*, convert_nodes: bool = True, convert_tensors: Union[bool, Literal['numpy']] = False, tolist_first: bool = False, as_linked_list: bool = False) → List[Any]¶

返回 tensordict 的嵌套列表表示。

如果 tensordict 没有批次维度，则此方法返回单个列表或字典。否则，它返回一个嵌套列表，其中每个内部列表代表一个批次维度。

参数:

convert_nodes (bool) – 如果为 True，则叶子节点将被转换为字典。否则，它们将被作为值列表返回。默认值：True。
convert_tensors (bool, "numpy") – 如果为 True，则在创建字典时，张量将被转换为列表。如果为“numpy”，则张量将被转换为 numpy 数组。否则，它们将保持为张量。默认值：False。
tolist_first (bool) – 如果为 True，则当 tensordict 具有批处理维度时，它将首先转换为列表。默认值：False。
as_linked_list (bool) – 如果为 True，则列表将被转换为 tensordict.utils.LinkedList，它将在修改列表时自动更新 tensordict。默认值：False。

返回:

tensordict 的嵌套列表表示。

示例

>>> import torch
>>> from tensordict import TensorDict
>>>
>>> td = TensorDict(
...     a=torch.arange(24).view(2, 3, 4),
...     b=TensorDict(c=torch.arange(12).reshape(2, 3, 2), batch_size=(2, 3, 2)),
...     batch_size=(2, 3)
... )
>>> print(td.tolist(tolist_first=True))
[[{'a': tensor([0, 1, 2, 3]), 'b': [{'c': tensor(0)}, {'c': tensor(1)}]}, {'a': tensor([4, 5, 6, 7]), 'b': [{'c': tensor(2)}, {'c': tensor(3)}]}, {'a': tensor([ 8,  9, 10, 11]), 'b': [{'c': tensor(4)}, {'c': tensor(5)}]}], [{'a': tensor([12, 13, 14, 15]), 'b': [{'c': tensor(6)}, {'c': tensor(7)}]}, {'a': tensor([16, 17, 18, 19]), 'b': [{'c': tensor(8)}, {'c': tensor(9)}]}, {'a': tensor([20, 21, 22, 23]), 'b': [{'c': tensor(10)}, {'c': tensor(11)}]}]]
>>> print(td.tolist(tolist_first=False))
[[{'a': tensor([0, 1, 2, 3]), 'b': {'c': tensor([0, 1])}}, {'a': tensor([4, 5, 6, 7]), 'b': {'c': tensor([2, 3])}}, {'a': tensor([ 8,  9, 10, 11]), 'b': {'c': tensor([4, 5])}}], [{'a': tensor([12, 13, 14, 15]), 'b': {'c': tensor([6, 7])}}, {'a': tensor([16, 17, 18, 19]), 'b': {'c': tensor([8, 9])}}, {'a': tensor([20, 21, 22, 23]), 'b': {'c': tensor([10, 11])}}]]
>>> print(td.tolist(convert_tensors=False))
[[{'a': [0, 1, 2, 3], 'b': [{'c': 0}, {'c': 1}]}, {'a': [4, 5, 6, 7], 'b': [{'c': 2}, {'c': 3}]}, {'a': [8, 9, 10, 11], 'b': [{'c': 4}, {'c': 5}]}], [{'a': [12, 13, 14, 15], 'b': [{'c': 6}, {'c': 7}]}, {'a': [16, 17, 18, 19], 'b': [{'c': 8}, {'c': 9}]}, {'a': [20, 21, 22, 23], 'b': [{'c': 10}, {'c': 11}]}]]
>>> print(td.tolist(convert_nodes=False))
[[[tensor([0, 1, 2, 3]), TensorDict(
    fields={
        c: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int64, is_shared=False)},
    batch_size=torch.Size([2]),
    device=None,
    is_shared=False)], [tensor([4, 5, 6, 7]), TensorDict(
    fields={
        c: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int64, is_shared=False)},
    batch_size=torch.Size([2]),
    device=None,
    is_shared=False)], [tensor([ 8,  9, 10, 11]), TensorDict(
    fields={
        c: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int64, is_shared=False)},
    batch_size=torch.Size([2]),
    device=None,
    is_shared=False)]], [[tensor([12, 13, 14, 15]), TensorDict(
    fields={
        c: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int64, is_shared=False)},
    batch_size=torch.Size([2]),
    device=None,
    is_shared=False)], [tensor([16, 17, 18, 19]), TensorDict(
    fields={
        c: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int64, is_shared=False)},
    batch_size=torch.Size([2]),
    device=None,
    is_shared=False)], [tensor([20, 21, 22, 23]), TensorDict(
    fields={
        c: Tensor(shape=torch.Size([2]), device=cpu, dtype=torch.int64, is_shared=False)},
    batch_size=torch.Size([2]),
    device=None,
    is_shared=False)]]]

transpose(dim0, dim1)¶

返回输入张量的转置张量。给定的维度 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() → Any¶: 计算 TensorDict 每个元素的 trunc() 值。

trunc_() → Any¶: 原地计算 TensorDict 每个元素的 trunc() 值。

type(dst_type: dtype) → Any¶

将所有张量转换为 dst_type。

参数:: dst_type (type or string) – 所需类型

uint16() → Any¶: 将所有张量转换为 torch.uint16。

uint32() → Any¶: 将所有张量转换为 torch.uint32。

uint64() → Any¶: 将所有张量转换为 torch.uint64。

uint8() → Any¶: 将所有张量转换为 torch.uint8。

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

返回一个已解绑的 tensordict 索引元组，沿指定的维度。

示例

>>> td = TensorDict({
...     'x': torch.arange(12).reshape(3, 4),
... }, batch_size=[3, 4])
>>> td0, td1, td2 = td.unbind(0)
>>> td0['x']
tensor([0, 1, 2, 3])
>>> td1['x']
tensor([4, 5, 6, 7])

unflatten(dim, unflattened_size)¶

展平 tensordict 的一个维度，将其扩展为所需的形状。

参数:

dim (int) – 指定要展平的输入张量的维度。
unflattened_size (shape) – 是 tensordict 展平维度的新的形状。

示例

>>> td = TensorDict({
...     "a": torch.arange(60).view(3, 4, 5),
...     "b": torch.arange(12).view(3, 4)},
...     batch_size=[3, 4])
>>> td_flat = td.flatten(0, 1)
>>> td_unflat = td_flat.unflatten(0, [3, 4])
>>> assert (td == td_unflat).all()

unflatten_keys(separator: str = '.', inplace: bool = False) → Any¶

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

TensorDict 类型将被丢失，结果将是简单的 TensorDict 实例。嵌套 tensordicts 的元数据将从根推断：数据树中的所有实例将共享相同的批次大小、维度名称和设备。

参数:

separator (str, optional) – 嵌套项之间的分隔符。
inplace (bool, optional) – 如果为 True，则结果 tensordict 将与调用它的 tensordict 具有相同的身份。默认为 False。

示例

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

此方法与 unflatten_keys() 特别适用于处理 state-dicts，因为它们可以无缝地将扁平字典转换为模仿模型结构的结构。

示例

>>> 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_() → Any¶

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

可用作装饰器。

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

unsqueeze(dim: int) → Self¶

为介于 -td.batch_dims 和 td.batch_dims 之间的维度扩展所有张量，并将它们返回到一个新的 tensordict 中。

参数:: dim (int) – 要扩展的维度

示例

>>> td = TensorDict({
...     'x': torch.arange(24).reshape(3, 4, 2),
... }, batch_size=[3, 4])
>>> td = td.unsqueeze(-2)
>>> td.shape
torch.Size([3, 1, 4])
>>> td.get("x").shape
torch.Size([3, 1, 4, 2])

此操作也可以用作上下文管理器。对原始 tensordict 的更改将是就地进行的，即原始张量的内容不会被修改。这也假定 tensordict 未被锁定（否则，需要解锁 tensordict）。

>>> td = TensorDict({
...     'x': torch.arange(24).reshape(3, 4, 2),
... }, batch_size=[3, 4])
>>> with td.unsqueeze(-2) as tds:
...     tds.set("y", torch.zeros(3, 1, 4))
>>> assert td.get("y").shape == [3, 4]

update(input_dict_or_td: Union[dict[str, Union[torch.Tensor, tensordict._tensorcollection.TensorCollection]], T], clone: bool = False, inplace: bool = False, *, non_blocking: bool = False, keys_to_update: Optional[Sequence[NestedKey]] = None, is_leaf: Optional[Callable[[Type], bool]] = None, update_batch_size: bool = False, ignore_lock: bool = False) → Any¶

使用来自字典或其他 TensorDict 的值来更新 TensorDict。

警告

update 在 try/except 块中调用时会损坏数据。不要在这些块中使用此方法来尝试捕获和修补执行过程中发生的错误。

参数:

input_dict_or_td (TensorDictBase 或 dict) – 要写入 self 的输入数据。
clone (bool, 可选) – 在设置输入（ tensor）字典中的张量之前是否克隆它们。默认为 False。
inplace (bool, 可选) – 如果为 True 并且某个键与 tensordict 中的现有键匹配，则更新将对该键值对进行原地更新。如果找不到该条目，则会将其添加。默认为 False。

关键字参数:

keys_to_update (NestedKeys 序列, 可选) – 如果提供，则仅更新 key_to_update 中的键列表。这旨在避免调用 data_dest.update(*keys_to_update))。
non_blocking (bool, 可选) – 如果为 True 且此复制在不同设备之间进行，则复制可能与主机异步发生。
is_leaf (Callable[[Type], bool], 可选) –
一个指示对象类型是应被视为叶子节点并被交换，还是应被视为张量集合的可调用对象。

另请参阅

is_leaf_nontensor() 和 default_is_leaf()。
update_batch_size (bool, 可选) –
如果为 True，则 update 将尝试更新目标（self）的 batch_size，如果它与源的 batch_size 不匹配。默认为 False。

注意

在 batch_size 不匹配的情况下，LazyStackTensorDict 实例将被清空其内容，并将使用源 tensordict 的副本重新填充容器。

注意

此参数假定 keys_to_update 为空，并且 inplace=False。如果目标（self）的键不是源键的子集，则会引发异常，因为 TensorDict 将无法推断如何处理额外的目标条目。
ignore_lock (bool, 可选) – 如果为 True，则无论其锁定状态如何，都可以更新任何 tensordict。默认为 False。

注意

当使用 N 个元素的 LazyStackedTensorDict 更新另一个 M 个元素的 LazyStackedTensorDict 时，沿堆栈维度，update 方法会将额外 tensordict 的副本附加到目标（self）的懒惰堆栈中。这使得用户可以依赖 update 来逐步增加懒惰堆栈。

返回:: 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: Union[dict[str, Union[torch.Tensor, tensordict._tensorcollection.TensorCollection]], T], clone: bool = False, *, non_blocking: bool = False, keys_to_update: Optional[Sequence[NestedKey]] = None) → Any¶

使用来自字典或其他 TensorDict 的值原地更新 TensorDict。

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

参数:

input_dict_or_td (TensorDictBase 或 dict) – 要写入 self 的输入数据。
clone (bool, 可选) – 在设置输入（ tensor）字典中的张量之前是否克隆它们。默认为 False。

关键字参数:

keys_to_update (NestedKeys 序列, 可选) – 如果提供，则仅更新 key_to_update 中的键列表。这旨在避免调用 data_dest.update_(data_src.select(*keys_to_update))。
non_blocking (bool, 可选) – 如果为 True 且此复制在不同设备之间进行，则复制可能与主机异步发生。

返回:

self

示例

>>> a = torch.randn(3)
>>> b = torch.randn(3, 4)
>>> td = TensorDict({"a": a, "b": b}, batch_size=[3])
>>> other_td = TensorDict({"a": a*0, "b": b*0}, batch_size=[])
>>> td.update_(other_td)
>>> assert td['a'] is not other_td['a']
>>> assert (td['a'] == other_td['a']).all()
>>> assert (td['a'] == 0).all()

update_at_(input_dict_or_td: Union[dict[str, Union[torch.Tensor, tensordict._tensorcollection.TensorCollection]], T], idx: Union[None, int, slice, str, Tensor, List[Any], Tuple[Any, ...]], clone: bool = False, *, non_blocking: bool = False, keys_to_update: Optional[Sequence[NestedKey]] = None) → Any¶

使用来自字典或其他 TensorDict 的值，在指定的索引处原地更新 TensorDict。

与 TensorDict.update 不同，此函数将在键未知于 TensorDict 时抛出错误。

参数:

input_dict_or_td (TensorDictBase 或 dict) – 要写入 self 的输入数据。
idx (int, torch.Tensor, 可迭代对象, slice) – 更新应发生的 tensordict 的索引。
clone (bool, 可选) – 在设置输入（ tensor）字典中的张量之前是否克隆它们。默认为 False。

关键字参数:

keys_to_update (NestedKeys序列, 可选) – 如果提供，将仅更新 key_to_update 中的键列表。
non_blocking (bool, 可选) – 如果为 True 且此复制在不同设备之间进行，则复制可能与主机异步发生。

返回:

self

示例

>>> td = TensorDict({
...     'a': torch.zeros(3, 4, 5),
...     'b': torch.zeros(3, 4, 10)}, batch_size=[3, 4])
>>> td.update_at_(
...     TensorDict({
...         'a': torch.ones(1, 4, 5),
...         'b': torch.ones(1, 4, 10)}, batch_size=[1, 4]),
...    slice(1, 2))
TensorDict(
    fields={
        a: Tensor(torch.Size([3, 4, 5]), dtype=torch.float32),
        b: Tensor(torch.Size([3, 4, 10]), dtype=torch.float32)},
    batch_size=torch.Size([3, 4]),
    device=None,
    is_shared=False)
>>> assert (td[1] == 1).all()

values(include_nested: bool = False, leaves_only: bool = False, is_leaf=None, *, sort: bool = False) → Iterator[Union[Tensor, TensorCollection]]¶

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

参数:

include_nested (bool, 可选) – 如果为 `True`，则返回嵌套值。默认为 `False`。
leaves_only (bool, 可选) – 如果为 `False`，则仅返回叶子节点。默认为 `False`。
is_leaf (callable, optional) –
一个作用于类类型的可调用对象，返回一个布尔值，指示该类是否应被视为叶子节点。

注意

is_leaf 的目的是不是阻止递归调用嵌套的 tensordicts，而是为过滤标记某些类型，当 `leaves_only=True` 时。即使 `is_leaf(cls)` 返回 `True`，如果 `include_nested=True`，嵌套 tensordict 的结构仍将被遍历。换句话说，`is_leaf` 不控制递归深度，而是提供了一种当 `leaves_only=True` 时从结果中过滤掉某些类型的方法。这意味着树中的一个节点既可以是叶节点，也可以是具有子节点的节点。实际上，`is_leaf` 的默认值不包含 tensordict 和 tensorclass 实例作为叶节点。

另请参阅

is_leaf_nontensor() 和 default_is_leaf()。

关键字参数:

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

var(dim: Union[int, Tuple[int], Literal['feature']] = NO_DEFAULT, keepdim: bool = NO_DEFAULT, *, correction: int = 1) → Any¶

var(dim: Union[int, Tuple[int], Literal['feature']] = NO_DEFAULT, keepdim: bool = NO_DEFAULT, *, correction: int = 1, reduce: bool) → Union[Any, Tensor]

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

参数:

dim (int, int 元组, 可选) – 如果为 None，则返回一个无量纲的 tensordict，其中包含所有叶子的总和值（如果可以计算）。如果为整数或整数元组，则当此维度与 tensordict 形状兼容时，将在指定的维度上调用 var。当前只允许 “feature” 字符串。使用 dim=”feature” 将在所有特征维度上执行归约。如果 reduce=True，则将返回形状与 TensorDict 的 batch_size 相同的张量。否则，将返回一个结构与 self 相同的、特征维度已归约的新 tensordict。
keepdim (bool) – 输出张量是否保留维度。

关键字参数:

correction (int) – 样本大小与样本自由度之间的差值。默认为贝塞尔校正，correction=1。
reduce (bool, optional) – 如果为 True，则会跨所有 TensorDict 值进行归约，并返回一个单一的归约张量。默认为 False。

示例

>>> from tensordict import TensorDict
>>> import torch
>>> td = TensorDict(
...     a=torch.randn(3, 4, 5),
...     b=TensorDict(
...         c=torch.randn(3, 4, 5, 6),
...         d=torch.randn(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.var(dim=0)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([4, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([4]),
    device=None,
    is_shared=False)
>>> td.var()
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([]), 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)
>>> td.var(reduce=True)
tensor(1.0006)
>>> td.var(dim="feature")
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False),
        b: TensorDict(
            fields={
                c: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False),
                d: Tensor(shape=torch.Size([3, 4, 5]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3, 4, 5]),
            device=None,
            is_shared=False)},
    batch_size=torch.Size([3, 4]),
    device=None,
    is_shared=False)
>>> td = TensorDict(
...     a=torch.ones(3, 4, 5),
...     b=TensorDict(
...         c=torch.ones(3, 4, 5),
...         d=torch.ones(3, 4, 5),
...         batch_size=(3, 4, 5),
...     ),
...     batch_size=(3, 4)
... )
>>> td.var(reduce=True, dim="feature")
tensor([[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]])
>>> td.var(reduce=True, dim=0)
tensor([[0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.]])

view(*shape: int)¶

view(dtype)

view(shape: Size)

返回一个 tensordict，其中包含与 tensordict batch_size 兼容的新形状的张量视图。

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

参数:

*shape (int) – 结果 TensorDict 的新形状。
dtype (torch.dtype) – 或者，用于表示张量内容的 dtype。
size – 可迭代对象

关键字参数:

batch_size (torch.Size, 可选) – 如果提供了 dtype，则可以使用此关键字参数重置 batch_size。如果 view 是使用形状调用的，则此参数无效。

返回:

具有所需 batch_size 的新 tensordict。

示例

>>> td = TensorDict(source={'a': torch.zeros(3,4,5),
...    'b': torch.zeros(3,4,10,1)}, batch_size=torch.Size([3, 4]))
>>> td_view = td.view(12)
>>> print(td_view.get("a").shape)  # torch.Size([12, 5])
>>> print(td_view.get("b").shape)  # torch.Size([12, 10, 1])
>>> td_view = td.view(-1, 4, 3)
>>> print(td_view.get("a").shape)  # torch.Size([1, 4, 3, 5])
>>> print(td_view.get("b").shape)  # torch.Size([1, 4, 3, 10, 1])

where(condition: Tensor, other: torch.Tensor | tensordict.base.TensorDictBase, *, out: Optional[TensorDictBase] = None, pad: Optional[Union[int, bool]] = None, update_batch_size: bool = False) → Any¶

根据 condition 从 self 或 other 中选择元素，并返回一个 TensorDict。

参数:

condition (BoolTensor) – 当 True (非零) 时，返回 self，否则返回 other。
other (TensorDictBase 或 Scalar) – 当 condition 为 False 时的索引处的值（如果 other 是标量）或值。

关键字参数:

out (TensorDictBase, 可选) – 输出的 TensorDictBase 实例。
pad (标量, 可选) – 如果提供，源或目标 tensordict 中缺失的键将写入为 torch.where(mask, self, pad) 或 torch.where(mask, pad, other)。默认为 None，即不允许缺失的键。
update_batch_size (bool, 可选) – 如果为 True 且提供了 out，则输出的 batch size 将更新为与 condition 的 batch size 匹配。默认为 False。

zero_() → Any¶: 原地将 tensordict 中的所有张量清零。

zero_grad(set_to_none: bool = True) → Any¶

递归地将 TensorDict 的所有梯度清零。

参数:: set_to_none (bool, 可选) – 如果为 True，则 tensor.grad 将为 None，否则为 0。默认为 True。

TensorDictBase¶

文档

教程

资源