Spatial Queries

This module provides spatial acceleration structures for efficient geometric queries on large meshes.

The BVH (Bounding Volume Hierarchy) is an axis-aligned bounding box tree built over mesh cells. It accelerates two key operations:

• Point containment: given a set of query points, find which mesh cell (if any) contains each point

• Nearest-cell search: find the closest cell to each query point

The BVH is used internally by the sampling module (sample_data_at_points(), find_containing_cells()) to avoid brute-force search over all cells.

import torch
from physicsnemo.mesh.spatial import BVH
from physicsnemo.mesh.primitives.surfaces import sphere_icosahedral

mesh = sphere_icosahedral.load(subdivisions=3)
bvh = BVH.from_mesh(mesh)

query_points = torch.randn(1000, 3)
candidate_cells = bvh.find_candidate_cells(query_points)

Signed Distance Field

signed_distance_field() computes the signed distance from a set of query points to a triangle surface mesh, together with the closest point on the surface for each query. It is a Mesh-typed wrapper around the Warp-backed physicsnemo.nn.functional.signed_distance_field() op, which runs NVIDIA Warp mesh queries on CPU and CUDA.

The sign is determined by one of two methods, selected with use_sign_winding_number:

• False (default): the angle-weighted pseudo-normal of the closest mesh feature (wp.mesh_query_point_sign_normal). This is fast and robust for watertight meshes.

• True: the generalized winding number (wp.mesh_query_point_sign_winding_number). This is robust for non-watertight / self-intersecting (“soup”) geometry.

import torch
from physicsnemo.mesh import Mesh
from physicsnemo.mesh.spatial import signed_distance_field

# A triangle surface mesh: (n_vertices, 3) coords + (n_faces, 3) connectivity.
mesh = Mesh(
    points=torch.randn(500, 3),
    cells=torch.randint(0, 500, (1000, 3)),
)

query = torch.randn(10000, 3)
sdf, hit_points, hit_faces = signed_distance_field(
    mesh, query, use_sign_winding_number=True
)
# sdf:        (10000,)   signed distances
# hit_points: (10000, 3) closest surface points
# hit_faces:  (10000,)   nearest-face index into mesh.cells

API Reference

Spatial acceleration structures for efficient queries on large meshes.

This module provides data structures and algorithms for fast spatial queries: - BVH (Bounding Volume Hierarchy) for point-in-cell queries - ClusterTree for dual-tree Barnes-Hut acceleration of kernel/attention operators - Signed distance field (signed_distance_field()) over a triangle

surface mesh, backed by NVIDIA Warp mesh queries (CPU and CUDA)

class physicsnemo.mesh.spatial.BVH(

node_aabb_min: jaxtyping.Float[Tensor, 'n_nodes n_spatial_dims'],

node_aabb_max: jaxtyping.Float[Tensor, 'n_nodes n_spatial_dims'],

node_left_child: jaxtyping.Int[Tensor, 'n_nodes'],

node_right_child: jaxtyping.Int[Tensor, 'n_nodes'],

leaf_start: jaxtyping.Int[Tensor, 'n_nodes'],

leaf_count: jaxtyping.Int[Tensor, 'n_nodes'],

sorted_cell_order: jaxtyping.Int[Tensor, 'n_cells'],

*,

batch_size,

device=None,

names=None,

)

Bases: object

property device: device

Device where BVH tensors are stored.

dumps(

prefix: str | None = None,

copy_existing: bool = False,

*,

num_threads: int = 0,

return_early: bool = False,

share_non_tensor: bool = False,

robust_key: bool | None = True,

archive: bool | None = None,

compression: str | int | None = None,

) → Any

Saves the tensordict to disk.

This function is a proxy to memmap().

classmethod fields()

Return a tuple describing the fields of this dataclass.

Accepts a dataclass or an instance of one. Tuple elements are of type Field.

find_candidate_cells(

query_points: Float[Tensor, 'n_queries n_spatial_dims'],

max_candidates_per_point: int | None = 32,

aabb_tolerance: float = 1e-06,

) → Adjacency

Find candidate cells that might contain each query point.

Uses batched iterative BVH traversal where all queries are processed simultaneously in a vectorized manner.

Parameters:

• query_points (torch.Tensor) – Points to query, shape (n_queries, n_spatial_dims).

• max_candidates_per_point (int | None, optional) – Maximum number of candidate cells to return per query point. Prevents memory explosion for degenerate cases. If None, no limit is applied.

• aabb_tolerance (float, optional) – Tolerance for AABB intersection test. Important for degenerate cells (e.g., cells with duplicate vertices).

Returns:

Adjacency object where candidates for query i are at result.indices[result.offsets[i]:result.offsets[i+1]]. Use result.to_list() for a list-of-tensors representation.

Return type:

Adjacency

Notes

Complexity is \(O(M \log N)\) where \(M\) = queries and \(N\) = cells. All AABB tests and tree operations are fully vectorized across queries - there are no Python-level loops over individual query points. The outer loop runs once per tree level (\(O(\log N)\) iterations).

from_csv(

*,

auto_batch_size: bool = False,

batch_dims: int | None = None,

device: device | None = None,

batch_size: Size | None = None,

separator: str | None = None,

dtype: dtype | None = None,

**kwargs,

) → Any

Creates a TensorDict from a CSV file.

Requires either pandas or pyarrow to be installed.

Parameters:

path (str or Path) – Path to the CSV file.

Keyword Arguments:

• auto_batch_size (bool, optional) – If True, the batch size will be computed automatically. Defaults to False.

• batch_dims (int, optional) – If auto_batch_size is True, defines how many dimensions the output tensordict should have. Defaults to None.

• device (torch.device, optional) – The device for tensor data. Defaults to None.

• batch_size (torch.Size, optional) – The batch size. Defaults to [num_rows].

• separator (str, optional) – If provided, column names are split on this separator to create nested TensorDicts. Defaults to None.

• dtype (torch.dtype, optional) – If provided, all numeric columns are cast to this dtype. Defaults to None.

• **kwargs – Additional keyword arguments forwarded to the CSV reader (pandas.read_csv or pyarrow.csv.read_csv).

Returns:

A TensorDict representation of the CSV data.

Examples

>>> td = TensorDict.from_csv("data.csv")
>>> td = TensorDict.from_csv("data.csv", separator=".", dtype=torch.float32)

from_json(

*,

auto_batch_size: bool = False,

batch_dims: int | None = None,

device: device | None = None,

batch_size: Size | None = None,

separator: str | None = None,

dtype: dtype | None = None,

lines: bool = False,

**kwargs,

) → Any

Creates a TensorDict from a JSON file.

Supports both standard JSON (array of records) and JSON Lines format. For nested JSON objects, use from_dict() instead.

Requires pandas for best results. Falls back to stdlib json for simple cases.

Parameters:

path (str or Path) – Path to the JSON file.

Keyword Arguments:

• auto_batch_size (bool, optional) – If True, the batch size will be computed automatically. Defaults to False.

• batch_dims (int, optional) – If auto_batch_size is True, defines how many dimensions the output tensordict should have. Defaults to None.

• device (torch.device, optional) – The device for tensor data. Defaults to None.

• batch_size (torch.Size, optional) – The batch size. Defaults to [num_rows].

• separator (str, optional) – If provided, column names are split on this separator to create nested TensorDicts. Defaults to None.

• dtype (torch.dtype, optional) – If provided, all numeric columns are cast to this dtype. Defaults to None.

• lines (bool, optional) – If True, reads the file as JSON Lines (one JSON object per line). Defaults to False.

• **kwargs – Additional keyword arguments forwarded to the JSON reader.

Returns:

A TensorDict representation of the JSON data.

Examples

>>> td = TensorDict.from_json("data.json")
>>> td = TensorDict.from_json("data.jsonl", lines=True)

classmethod from_mesh(mesh: Mesh, leaf_size: int = 1) → BVH

Construct a BVH from a mesh using morton-code LBVH.

Cells are sorted by the morton code of their centroids, then the tree is built top-down by recursively splitting sorted segments at their midpoints. AABBs are computed in two phases: leaf AABBs via segmented reduction over cell bounds, internal AABBs via a bottom-up pass from leaves to root. The entire construction runs in O(log N) Python-level iterations.

Parameters:

• mesh (Mesh) – The mesh to build the BVH for.

• leaf_size (int, optional) – Maximum number of cells per leaf node. The default of 1 minimizes candidate cells per query hit but maximizes node count (2 * n_cells - 1 nodes) and tree depth. Larger values reduce build time and memory at the cost of more candidate cells per query hit.

Returns:

Constructed BVH ready for queries.

Return type:

BVH

Raises:

ValueError – If leaf_size < 1.

from_pandas(

*,

auto_batch_size: bool = False,

batch_dims: int | None = None,

device: device | None = None,

batch_size: Size | None = None,

separator: str | None = None,

dtype: dtype | None = None,

) → Any

Converts a pandas DataFrame to a TensorDict.

Numeric columns become tensors, string/object columns become NonTensorData.

Parameters:

dataframe (pd.DataFrame) – The pandas DataFrame to convert.

Keyword Arguments:

• auto_batch_size (bool, optional) – If True, the batch size will be computed automatically. Defaults to False.

• batch_dims (int, optional) – If auto_batch_size is True, defines how many dimensions the output tensordict should have. Defaults to None.

• device (torch.device, optional) – The device for tensor data. Defaults to None.

• batch_size (torch.Size, optional) – The batch size. Defaults to [num_rows].

• separator (str, optional) – If provided, column names are split on this separator to create nested TensorDicts. For example, with separator=".", a column "obs.x" becomes td["obs", "x"]. Defaults to None.

• dtype (torch.dtype, optional) – If provided, all numeric columns are cast to this dtype. Defaults to None.

Returns:

A TensorDict representation of the DataFrame.

Examples

>>> import pandas as pd
>>> df = pd.DataFrame({"a": [1, 2, 3], "b": [4.0, 5.0, 6.0]})
>>> td = TensorDict.from_pandas(df)
>>> print(td)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.int64, is_shared=False),
        b: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float64, is_shared=False)},
    batch_size=torch.Size([3]),
    device=None,
    is_shared=False)

from_parquet(

*,

auto_batch_size: bool = False,

batch_dims: int | None = None,

device: device | None = None,

batch_size: Size | None = None,

separator: str | None = None,

dtype: dtype | None = None,

columns: list[str] | None = None,

**kwargs,

) → Any

Creates a TensorDict from a Parquet file.

Requires either pyarrow or pandas to be installed. Prefers pyarrow when available for better performance.

Parameters:

path (str or Path) – Path to the Parquet file.

Keyword Arguments:

• auto_batch_size (bool, optional) – If True, the batch size will be computed automatically. Defaults to False.

• batch_dims (int, optional) – If auto_batch_size is True, defines how many dimensions the output tensordict should have. Defaults to None.

• device (torch.device, optional) – The device for tensor data. Defaults to None.

• batch_size (torch.Size, optional) – The batch size. Defaults to [num_rows].

• separator (str, optional) – If provided, column names are split on this separator to create nested TensorDicts. Defaults to None.

• dtype (torch.dtype, optional) – If provided, all numeric columns are cast to this dtype. Defaults to None.

• columns (list of str, optional) – If provided, only read these columns from the file. Defaults to None (all columns).

• **kwargs – Additional keyword arguments forwarded to the Parquet reader.

Returns:

A TensorDict representation of the Parquet data.

Examples

>>> td = TensorDict.from_parquet("data.parquet")
>>> td = TensorDict.from_parquet("data.parquet", columns=["obs", "reward"])

from_schema(

*,

batch_size: Sequence[int] | Size | None = None,

storage: str | None = None,

device=None,

**kwargs,

) → TensorDictBase

Pre-allocate a zero-filled TensorDict from a schema.

Creates a TensorDictBase whose storage backend is selected by storage. Each entry in schema maps a field name to an (element_shape, dtype) pair; the full stored shape is [*batch_size, *element_shape].

Parameters:

schema – Mapping from field name to (element_shape, dtype). element_shape is the per-element shape (excluding batch_size).

Keyword Arguments:

• batch_size – Overall batch dimensions prepended to every element shape. Defaults to ().

• storage (str or None) –

  Backend selector:

  • None – plain TensorDict with regular tensors.

  • "memmap" – memory-mapped tensors on disk. Pass prefix=<dir> in kwargs.

  • "h5" – HDF5 via PersistentTensorDict. Pass filename=<path> in kwargs.

  • "zarr" – zarr (requires zarr>=3.0) via PersistentTensorDict. Pass filename=<path or store> in kwargs.

  • "shared" – CPU shared-memory tensors.

  • "redis" / "dragonfly" – delegates to TensorDictStore.from_schema().

• device – Device for the resulting tensors (ignored by some backends).

• **kwargs – Backend-specific arguments forwarded to the underlying constructor (e.g. prefix for memmap, filename for h5, host/port for redis).

Returns:

A new TensorDictBase subclass instance with pre-allocated (zero-filled) keys.

Examples

>>> td = TensorDict.from_schema(
...     {"obs": ([84, 84, 3], torch.uint8),
...      "reward": ([], torch.float32)},
...     batch_size=[1000],
... )
>>> td["obs"].shape
torch.Size([1000, 84, 84, 3])

>>> import tempfile
>>> with tempfile.TemporaryDirectory() as d:
...     td_mm = TensorDict.from_schema(
...         {"obs": ([4], torch.float32)},
...         batch_size=[8],
...         storage="memmap",
...         prefix=d,
...     )
...     assert td_mm.is_memmap()

classmethod from_tensordict(

tensordict: TensorDictBase,

non_tensordict: dict | None = None,

safe: bool = True,

) → Any

Tensor class wrapper to instantiate a new tensor class object.

Parameters:

• tensordict (TensorDictBase) – Dictionary of tensor types

• non_tensordict (dict) – Dictionary with non-tensor and nested tensor class objects

• safe (bool) – Whether to raise an error if the tensordict is not a TensorDictBase instance

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

Gets the value stored with the input key.

Parameters:

• key (str, tuple of str) – key to be queried. If tuple of str it is equivalent to chained calls of getattr.

• default – default value if the key is not found in the tensorclass.

Returns:

value stored with the input key

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

Loads a tensordict from disk.

This class method is a proxy to load_memmap().

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

Loads a tensordict from disk within the current tensordict.

This class method is a proxy to load_memmap_().

classmethod load_memmap(

prefix: str | Path,

device: device | None = None,

non_blocking: bool = False,

*,

out: TensorDictBase | None = None,

robust_key: bool | None = True,

subpath: NestedKey | None = None,

mode: str = 'r',

num_threads: int = 0,

allow_pickle: bool | None = None,

) → Any

Loads a memory-mapped tensordict from disk.

Parameters:

• prefix (str or Path to folder) – the path to the folder where the saved tensordict should be fetched, or the path to a memmap archive file written through save(..., archive=True) / a ".tdz" prefix (or packed with pack_memmap()). Archives are memory-mapped once and every leaf is exposed as a zero-copy view into the mapping: only the pages of the leaves that are actually accessed are read from disk. Unlike directory-backed tensordicts, in-place writes to the leaves of an archive-loaded tensordict do not propagate to the file.

• device (torch.device or equivalent, optional) – if provided, the data will be asynchronously cast to that device. Supports “meta” device, in which case the data isn’t loaded but a set of empty “meta” tensors are created. This is useful to get a sense of the total model size and structure without actually opening any file.

• non_blocking (bool, optional) – if True, synchronize won’t be called after loading tensors on device. Defaults to False.

• out (TensorDictBase, optional) – optional tensordict where the data should be written.

• robust_key (bool, optional) – if True (default), expects robust key encoding was used when saving and decodes filenames accordingly. If False, uses legacy behavior. If None, uses the default robust behavior.

• subpath (NestedKey or str path, optional) – the location of a nested tensordict to load, as a nested key (e.g. ("module", "0"), with arbitrary nesting allowed as usual) or as a "/"-separated string path (e.g. "module/0"). Only that subtree is loaded. Works both for directories (equivalent to appending the path to prefix) and for archives.

• mode (str, optional) – "r" (default) or "r+". Only relevant when loading an archive: with "r" the archive is mapped copy-on-write and in-place writes to the leaves stay in memory; with "r+" the mapping is shared and in-place writes propagate to the file, like directory-backed tensordicts. "r+" requires uncompressed, aligned tensor payloads (i.e. archives written by tensordict without compression) and is not available for nested-tensor leaves. In-place writes do not update the per-entry CRC-32 stored by the zip format; load_memmap() ignores checksums, but call refresh_archive_checksums() before handing a modified archive to tools that verify them (unzip, unpack_memmap(), …). Directory prefixes are always write-through and ignore this argument.

• num_threads (int, optional) – number of threads used to decompress the leaves of a compressed archive (deflate entries are inflated in parallel, which scales nearly linearly). Without compression, loading is a metadata-only operation and this argument has no effect. Defaults to 0 (sequential).

• allow_pickle (bool, optional) – whether pickled non-tensor fields may be loaded. Pickle can execute arbitrary code, so pass True only for data from a trusted source and False for untrusted data. During the 0.14 compatibility window, omitting this option loads pickle with a FutureWarning; the default will change to False in 0.15. Saves without a pickle sidecar do not require this option.

Examples

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

This method also allows loading nested tensordicts.

Examples

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

A tensordict can also be loaded on “meta” device or, alternatively, as a fake tensor.

Examples

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

state_dict: dict[str, Any],

strict=True,

assign=False,

from_flatten=None,

)

Loads a state_dict into the tensorclass.

Supports both the new format (logical keys with _metadata) and the legacy format (_tensordict/_non_tensordict wrapper keys).

memmap(

prefix: str | None = None,

copy_existing: bool = False,

*,

num_threads: int = 0,

return_early: bool = False,

share_non_tensor: bool = False,

existsok: bool = True,

robust_key: bool | None = True,

archive: bool | None = None,

compression: str | int | None = None,

) → Any

Writes all tensors onto a corresponding memory-mapped Tensor in a new tensordict.

Parameters:

• prefix (str) – directory prefix where the memory-mapped tensors will be stored. The directory tree structure will mimic the tensordict’s. If prefix ends with ".tdz" (or archive=True is passed), a single-file archive is written instead of a directory: a standard zip file whose entries replicate the memmap directory layout. See archive below.

• copy_existing (bool) – If False (default), an exception will be raised if an entry in the tensordict is already a tensor stored on disk with an associated file, but is not saved in the correct location according to prefix. If True, any existing Tensor will be copied to the new location.

Keyword Arguments:

• num_threads (int, optional) – the number of threads used to write the memmap tensors. Defaults to 0.

• return_early (bool, optional) – if True and num_threads>0, the method will return a future of the tensordict.

• share_non_tensor (bool, optional) – if True, the non-tensor data will be shared between the processes and writing operation (such as inplace update or set) on any of the workers within a single node will update the value on all other workers. If the number of non_tensor leaves is high (e.g., sharing large stacks of non-tensor data) this may result in OOM or similar errors. Defaults to False.

• existsok (bool, optional) – if False, an exception will be raised if a tensor already exists in the same path. Defaults to True.

• robust_key (bool, optional) – if True (default), uses robust key encoding that safely handles keys with path separators and special characters. If False, uses legacy behavior (keys used as-is). If None, uses the default robust behavior.

• archive (bool, optional) – if True, prefix designates a single file rather than a directory and the tensordict is written as a memmap archive: a zip file mirroring the memmap directory tree, with tensor payloads stored uncompressed and aligned so that load_memmap() can memory-map the file and expose every leaf as a zero-copy view. If None (default), archive mode is enabled when prefix ends with ".tdz". The result of load_memmap() on an archive behaves like the result of from_consolidated(): all leaves are views into a single storage, and in-place writes do not propagate to the file. Archives and memmap directories are mutually convertible with pack_memmap() / unpack_memmap() (or any zip tool). Note that archives are written sequentially (single data pass) and num_threads has no effect on them.

• compression (str or int, optional) – compression for archive entries ("stored", "deflate", "bzip2", "lzma" or a zipfile constant). Defaults to "stored" (uncompressed), which is what enables zero-copy loading. Compressed archives load correctly but leaves are decompressed in memory on access. Only valid in archive mode.

The TensorDict is then locked, meaning that any writing operations that isn’t in-place will throw an exception (eg, rename, set or remove an entry). Once the tensordict is unlocked, the memory-mapped attribute is turned to False, because cross-process identity is not guaranteed anymore.

Returns:

A new tensordict with the tensors stored on disk if return_early=False, otherwise a TensorDictFuture instance.

Note

Serialising in this fashion might be slow with deeply nested tensordicts, so it is not recommended to call this method inside a training loop.

memmap_(

prefix: str | None = None,

copy_existing: bool = False,

*,

num_threads: int = 0,

return_early: bool = False,

share_non_tensor: bool = False,

existsok: bool = True,

robust_key: bool | None = True,

) → Any

Writes all tensors onto a corresponding memory-mapped Tensor, in-place.

Parameters:

• prefix (str) – directory prefix where the memory-mapped tensors will be stored. The directory tree structure will mimic the tensordict’s.

• copy_existing (bool) – If False (default), an exception will be raised if an entry in the tensordict is already a tensor stored on disk with an associated file, but is not saved in the correct location according to prefix. If True, any existing Tensor will be copied to the new location.

Keyword Arguments:

• num_threads (int, optional) – the number of threads used to write the memmap tensors. Defaults to 0.

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

• share_non_tensor (bool, optional) – if True, the non-tensor data will be shared between the processes and writing operation (such as inplace update or set) on any of the workers within a single node will update the value on all other workers. If the number of non-tensor leaves is high (e.g., sharing large stacks of non-tensor data) this may result in OOM or similar errors. Defaults to False.

• existsok (bool, optional) – if False, an exception will be raised if a tensor already exists in the same path. Defaults to True.

• robust_key (bool, optional) – if True (default), uses robust key encoding that safely handles keys with path separators and special characters. If False, uses legacy behavior (keys used as-is). If None, uses the default robust behavior.

The TensorDict is then locked, meaning that any writing operations that isn’t in-place will throw an exception (eg, rename, set or remove an entry). Once the tensordict is unlocked, the memory-mapped attribute is turned to False, because cross-process identity is not guaranteed anymore.

Returns:

self if return_early=False, otherwise a TensorDictFuture instance.

Note

Serialising in this fashion might be slow with deeply nested tensordicts, so it is not recommended to call this method inside a training loop.

memmap_like(

prefix: str | None = None,

copy_existing: bool = False,

*,

existsok: bool = True,

num_threads: int = 0,

return_early: bool = False,

share_non_tensor: bool = False,

robust_key: bool | None = True,

archive: bool | None = None,

) → Any

Creates a contentless Memory-mapped tensordict with the same shapes as the original one.

Parameters:

• prefix (str) – directory prefix where the memory-mapped tensors will be stored. The directory tree structure will mimic the tensordict’s. If prefix ends with ".tdz" (or archive=True is passed), a preallocated single-file archive is created instead. See archive below.

• copy_existing (bool) – If False (default), an exception will be raised if an entry in the tensordict is already a tensor stored on disk with an associated file, but is not saved in the correct location according to prefix. If True, any existing Tensor will be copied to the new location.

Keyword Arguments:

• num_threads (int, optional) – the number of threads used to write the memmap tensors. Defaults to 0.

• return_early (bool, optional) – if True and num_threads>0, the method will return a future of the tensordict.

• share_non_tensor (bool, optional) – if True, the non-tensor data will be shared between the processes and writing operation (such as inplace update or set) on any of the workers within a single node will update the value on all other workers. If the number of non-tensor leaves is high (e.g., sharing large stacks of non-tensor data) this may result in OOM or similar errors. Defaults to False.

• existsok (bool, optional) – if False, an exception will be raised if a tensor already exists in the same path. Defaults to True.

• robust_key (bool, optional) – if True (default), uses robust key encoding that safely handles keys with path separators and special characters. If False, uses legacy behavior (keys used as-is). If None, uses the default robust behavior.

• archive (bool, optional) – if True, prefix designates a single file and a preallocated, zero-filled memmap archive is created and loaded back with load_memmap(prefix, mode="r+"): the returned tensordict writes through to the archive. If None (default), archive mode is enabled when prefix ends with ".tdz". In-place writes leave the zip per-entry checksums stale; call refresh_archive_checksums() before handing the archive to tools that verify them. Nested tensors are not supported in this mode.

The TensorDict is then locked, meaning that any writing operations that isn’t in-place will throw an exception (eg, rename, set or remove an entry). Once the tensordict is unlocked, the memory-mapped attribute is turned to False, because cross-process identity is not guaranteed anymore.

Returns:

A new TensorDict instance with data stored as memory-mapped tensors if return_early=False, otherwise a TensorDictFuture instance.

Note

This is the recommended method to write a set of large buffers on disk, as memmap_() will copy the information, which can be slow for large content.

Examples

>>> 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_(*, allow_pickle: bool | None = None)

Refreshes the content of the memory-mapped tensordict if it has a saved_path.

This method will raise an exception if no path is associated with it.

Parameters:

allow_pickle (bool, optional) – whether pickled non-tensor fields may be loaded. See load_memmap().

property n_nodes: int

Number of nodes in the BVH.

property n_spatial_dims: int

Dimensionality of the spatial space.

point_in_aabb(

points: Float[Tensor, 'n_points n_spatial_dims'],

aabb_min: Float[Tensor, 'n_boxes n_spatial_dims'],

aabb_max: Float[Tensor, 'n_boxes n_spatial_dims'],

) → Bool[Tensor, 'n_points n_boxes']

Test if points are inside axis-aligned bounding boxes.

Parameters:

• points (torch.Tensor) – Query points, shape (n_points, n_spatial_dims).

• aabb_min (torch.Tensor) – Minimum corners, shape (n_boxes, n_spatial_dims).

• aabb_max (torch.Tensor) – Maximum corners, shape (n_boxes, n_spatial_dims).

Returns:

Boolean tensor of shape (n_points, n_boxes) indicating containment.

Return type:

torch.Tensor

save(

prefix: str | None = None,

copy_existing: bool = False,

*,

num_threads: int = 0,

return_early: bool = False,

share_non_tensor: bool = False,

robust_key: bool | None = True,

archive: bool | None = None,

compression: str | int | None = None,

) → Any

Saves the tensordict to disk.

This function is a proxy to memmap().

select(

*keys,

inplace: bool = False,

strict: bool = True,

as_tensordict: bool = False,

)

TensorClass-specific select that supports as_tensordict.

set(

key: NestedKey,

value: Any,

inplace: bool = False,

non_blocking: bool = False,

)

Sets a new key-value pair.

Parameters:

• key (str, tuple of str) – name of the key to be set. If tuple of str it is equivalent to chained calls of getattr followed by a final setattr.

• value (Any) – value to be stored in the tensorclass

• inplace (bool, optional) – if True, set will tentatively try to update the value in-place. If False or if the key isn’t present, the value will be simply written at its destination.

Returns:

self

state_dict(

destination=None,

prefix='',

keep_vars=False,

flatten=True,

) → dict[str, Any]

Returns a state_dict with logical keys, matching TensorDictBase conventions.

Tensor fields appear as data keys. Non-tensor fields (strings, ints, etc.) and the tensorclass type are stored in _metadata. This replaces the legacy _tensordict/_non_tensordict wrapper format.

to_tensordict(

*,

retain_none: bool | None = None,

) → TensorDict

Convert the tensorclass into a regular TensorDict.

Makes a copy of all entries. Memmap and shared memory tensors are converted to regular tensors.

Parameters:

retain_none (bool) – if True, the None values will be written in the tensordict. Otherwise they will be discrarded. Default: True.

Returns:

A new TensorDict object containing the same values as the tensorclass.

unbind(dim: int)

Returns a tuple of indexed tensorclass instances unbound along the indicated dimension.

Resulting tensorclass instances will share the storage of the initial tensorclass instance.

class physicsnemo.mesh.spatial.ClusterTree(

node_aabb_min: torch.Tensor,

node_aabb_max: torch.Tensor,

node_diameter_sq: torch.Tensor,

node_left_child: torch.Tensor,

node_right_child: torch.Tensor,

leaf_start: torch.Tensor,

leaf_count: torch.Tensor,

node_range_start: torch.Tensor,

node_range_count: torch.Tensor,

node_total_area: torch.Tensor,

sorted_source_order: torch.Tensor,

source_points: torch.Tensor,

max_depth: torch.Tensor,

*,

batch_size,

device=None,

names=None,

)

Bases: object

compute_source_aggregates(

source_points: Float[Tensor, 'n_sources n_dims'],

areas: Float[Tensor, 'n_sources'],

source_data: TensorDict | None = None,

) → SourceAggregates

Compute per-node aggregate source data for far-field approximation.

Aggregates are area-weighted averages of source features within each node’s subtree. The areas passed to this call are authoritative for both the weighted sums and their normalization; they need not match the areas used to construct the tree. Per-source strengths are handled separately during kernel evaluation.

Parameters:

• source_points (Float[torch.Tensor, "n_sources n_dims"]) – Source coordinates, shape \((N, D)\).

• areas (Float[torch.Tensor, "n_sources"]) – Per-source area weights, shape \((N,)\).

• source_data (TensorDict or None) – Per-source features (normals, latents, etc.) with batch_size=(N,). None if no per-source features.

Returns:

Per-node aggregated centroids and source data.

Return type:

SourceAggregates

Notes

Tree topology depends on point positions, so callers may reuse a cached tree with different aggregation weights. node_total_area remains construction-time metadata; it is not the normalization for this call’s aggregates.

property device: device

Retrieves the device type of tensor class.

dumps(

prefix: str | None = None,

copy_existing: bool = False,

*,

num_threads: int = 0,

return_early: bool = False,

share_non_tensor: bool = False,

robust_key: bool | None = True,

archive: bool | None = None,

compression: str | int | None = None,

) → Any

Saves the tensordict to disk.

This function is a proxy to memmap().

classmethod fields()

Return a tuple describing the fields of this dataclass.

Accepts a dataclass or an instance of one. Tuple elements are of type Field.

find_dual_interaction_pairs(

target_tree: ClusterTree,

theta: float = 1.0,

*,

expand_far_targets: bool = False,

) → DualInteractionPlan

Find near-field and far-field pairs via dual-tree traversal.

Traverses both the source tree (self) and target_tree simultaneously. For well-separated node pairs, records a single far-field (target_node, source_node) entry - the kernel is evaluated ONCE at the node centroids and broadcast to all targets in the node. This reduces far-field kernel evaluations from O(N log N) to O(N).

Uses a combined AABB-distance opening criterion: (D_T + D_S) / r < theta, where D_T and D_S are the AABB diagonals and r is the minimum distance between the two AABBs. This accounts for approximation error on both the target and source sides.

Parameters:

• target_tree (ClusterTree) – Tree over target points. For self-interaction (communication layers), this is the same object as self.

• theta (float) – Barnes-Hut opening angle. Larger = more aggressive. theta = 0 forces all interactions to be exact.

• expand_far_targets (bool, optional, default=False) – If True, far-field node pairs are expanded to individual target points, converting (far, far) entries into (near, far) entries. This eliminates the target-side centroid approximation (and the blocky spatial artifacts it produces) at the cost of more kernel evaluations while preserving the source-side monopole speedup.

Returns:

Near-field individual pairs and far-field node-to-node pairs.

Return type:

DualInteractionPlan

from_csv(

*,

auto_batch_size: bool = False,

batch_dims: int | None = None,

device: device | None = None,

batch_size: Size | None = None,

separator: str | None = None,

dtype: dtype | None = None,

**kwargs,

) → Any

Creates a TensorDict from a CSV file.

Requires either pandas or pyarrow to be installed.

Parameters:

path (str or Path) – Path to the CSV file.

Keyword Arguments:

• auto_batch_size (bool, optional) – If True, the batch size will be computed automatically. Defaults to False.

• batch_dims (int, optional) – If auto_batch_size is True, defines how many dimensions the output tensordict should have. Defaults to None.

• device (torch.device, optional) – The device for tensor data. Defaults to None.

• batch_size (torch.Size, optional) – The batch size. Defaults to [num_rows].

• separator (str, optional) – If provided, column names are split on this separator to create nested TensorDicts. Defaults to None.

• dtype (torch.dtype, optional) – If provided, all numeric columns are cast to this dtype. Defaults to None.

• **kwargs – Additional keyword arguments forwarded to the CSV reader (pandas.read_csv or pyarrow.csv.read_csv).

Returns:

A TensorDict representation of the CSV data.

Examples

>>> td = TensorDict.from_csv("data.csv")
>>> td = TensorDict.from_csv("data.csv", separator=".", dtype=torch.float32)

from_json(

*,

auto_batch_size: bool = False,

batch_dims: int | None = None,

device: device | None = None,

batch_size: Size | None = None,

separator: str | None = None,

dtype: dtype | None = None,

lines: bool = False,

**kwargs,

) → Any

Creates a TensorDict from a JSON file.

Supports both standard JSON (array of records) and JSON Lines format. For nested JSON objects, use from_dict() instead.

Requires pandas for best results. Falls back to stdlib json for simple cases.

Parameters:

path (str or Path) – Path to the JSON file.

Keyword Arguments:

• auto_batch_size (bool, optional) – If True, the batch size will be computed automatically. Defaults to False.

• batch_dims (int, optional) – If auto_batch_size is True, defines how many dimensions the output tensordict should have. Defaults to None.

• device (torch.device, optional) – The device for tensor data. Defaults to None.

• batch_size (torch.Size, optional) – The batch size. Defaults to [num_rows].

• separator (str, optional) – If provided, column names are split on this separator to create nested TensorDicts. Defaults to None.

• dtype (torch.dtype, optional) – If provided, all numeric columns are cast to this dtype. Defaults to None.

• lines (bool, optional) – If True, reads the file as JSON Lines (one JSON object per line). Defaults to False.

• **kwargs – Additional keyword arguments forwarded to the JSON reader.

Returns:

A TensorDict representation of the JSON data.

Examples

>>> td = TensorDict.from_json("data.json")
>>> td = TensorDict.from_json("data.jsonl", lines=True)

from_pandas(

*,

auto_batch_size: bool = False,

batch_dims: int | None = None,

device: device | None = None,

batch_size: Size | None = None,

separator: str | None = None,

dtype: dtype | None = None,

) → Any

Converts a pandas DataFrame to a TensorDict.

Numeric columns become tensors, string/object columns become NonTensorData.

Parameters:

dataframe (pd.DataFrame) – The pandas DataFrame to convert.

Keyword Arguments:

• auto_batch_size (bool, optional) – If True, the batch size will be computed automatically. Defaults to False.

• batch_dims (int, optional) – If auto_batch_size is True, defines how many dimensions the output tensordict should have. Defaults to None.

• device (torch.device, optional) – The device for tensor data. Defaults to None.

• batch_size (torch.Size, optional) – The batch size. Defaults to [num_rows].

• separator (str, optional) – If provided, column names are split on this separator to create nested TensorDicts. For example, with separator=".", a column "obs.x" becomes td["obs", "x"]. Defaults to None.

• dtype (torch.dtype, optional) – If provided, all numeric columns are cast to this dtype. Defaults to None.

Returns:

A TensorDict representation of the DataFrame.

Examples

>>> import pandas as pd
>>> df = pd.DataFrame({"a": [1, 2, 3], "b": [4.0, 5.0, 6.0]})
>>> td = TensorDict.from_pandas(df)
>>> print(td)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.int64, is_shared=False),
        b: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float64, is_shared=False)},
    batch_size=torch.Size([3]),
    device=None,
    is_shared=False)

from_parquet(

*,

auto_batch_size: bool = False,

batch_dims: int | None = None,

device: device | None = None,

batch_size: Size | None = None,

separator: str | None = None,

dtype: dtype | None = None,

columns: list[str] | None = None,

**kwargs,

) → Any

Creates a TensorDict from a Parquet file.

Requires either pyarrow or pandas to be installed. Prefers pyarrow when available for better performance.

Parameters:

path (str or Path) – Path to the Parquet file.

Keyword Arguments:

• auto_batch_size (bool, optional) – If True, the batch size will be computed automatically. Defaults to False.

• batch_dims (int, optional) – If auto_batch_size is True, defines how many dimensions the output tensordict should have. Defaults to None.

• device (torch.device, optional) – The device for tensor data. Defaults to None.

• batch_size (torch.Size, optional) – The batch size. Defaults to [num_rows].

• separator (str, optional) – If provided, column names are split on this separator to create nested TensorDicts. Defaults to None.

• dtype (torch.dtype, optional) – If provided, all numeric columns are cast to this dtype. Defaults to None.

• columns (list of str, optional) – If provided, only read these columns from the file. Defaults to None (all columns).

• **kwargs – Additional keyword arguments forwarded to the Parquet reader.

Returns:

A TensorDict representation of the Parquet data.

Examples

>>> td = TensorDict.from_parquet("data.parquet")
>>> td = TensorDict.from_parquet("data.parquet", columns=["obs", "reward"])

classmethod from_points(

points: Float[Tensor, 'n_points n_dims'],

*,

leaf_size: int = 1,

areas: Float[Tensor, 'n_points'] | None = None,

) → ClusterTree

Build a cluster tree from a set of points via morton-code LBVH.

Parameters:

• points (Float[torch.Tensor, "n_points n_dims"]) – Source point coordinates, shape \((N, D)\).

• leaf_size (int) – Maximum sources per leaf node. Larger values produce shallower trees (fewer traversal iterations) at the cost of more exact near-field interactions per leaf hit.

• areas (Float[torch.Tensor, "n_points"] or None) – Per-source area weights used for aggregate computation. If None, all areas default to 1.

Returns:

Constructed tree ready for traversal and aggregate computation.

Return type:

ClusterTree

from_schema(

*,

batch_size: Sequence[int] | Size | None = None,

storage: str | None = None,

device=None,

**kwargs,

) → TensorDictBase

Pre-allocate a zero-filled TensorDict from a schema.

Creates a TensorDictBase whose storage backend is selected by storage. Each entry in schema maps a field name to an (element_shape, dtype) pair; the full stored shape is [*batch_size, *element_shape].

Parameters:

schema – Mapping from field name to (element_shape, dtype). element_shape is the per-element shape (excluding batch_size).

Keyword Arguments:

• batch_size – Overall batch dimensions prepended to every element shape. Defaults to ().

• storage (str or None) –

  Backend selector:

  • None – plain TensorDict with regular tensors.

  • "memmap" – memory-mapped tensors on disk. Pass prefix=<dir> in kwargs.

  • "h5" – HDF5 via PersistentTensorDict. Pass filename=<path> in kwargs.

  • "zarr" – zarr (requires zarr>=3.0) via PersistentTensorDict. Pass filename=<path or store> in kwargs.

  • "shared" – CPU shared-memory tensors.

  • "redis" / "dragonfly" – delegates to TensorDictStore.from_schema().

• device – Device for the resulting tensors (ignored by some backends).

• **kwargs – Backend-specific arguments forwarded to the underlying constructor (e.g. prefix for memmap, filename for h5, host/port for redis).

Returns:

A new TensorDictBase subclass instance with pre-allocated (zero-filled) keys.

Examples

>>> td = TensorDict.from_schema(
...     {"obs": ([84, 84, 3], torch.uint8),
...      "reward": ([], torch.float32)},
...     batch_size=[1000],
... )
>>> td["obs"].shape
torch.Size([1000, 84, 84, 3])

>>> import tempfile
>>> with tempfile.TemporaryDirectory() as d:
...     td_mm = TensorDict.from_schema(
...         {"obs": ([4], torch.float32)},
...         batch_size=[8],
...         storage="memmap",
...         prefix=d,
...     )
...     assert td_mm.is_memmap()

classmethod from_tensordict(

tensordict: TensorDictBase,

non_tensordict: dict | None = None,

safe: bool = True,

) → Any

Tensor class wrapper to instantiate a new tensor class object.

Parameters:

• tensordict (TensorDictBase) – Dictionary of tensor types

• non_tensordict (dict) – Dictionary with non-tensor and nested tensor class objects

• safe (bool) – Whether to raise an error if the tensordict is not a TensorDictBase instance

get(

key: NestedKey,

*args,

**kwargs,

)

Gets the value stored with the input key.

Parameters:

• key (str, tuple of str) – key to be queried. If tuple of str it is equivalent to chained calls of getattr.

• default – default value if the key is not found in the tensorclass.

Returns:

value stored with the input key

classmethod load(

prefix: str | Path,

*args,

**kwargs,

) → Any

Loads a tensordict from disk.

This class method is a proxy to load_memmap().

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

Loads a tensordict from disk within the current tensordict.

This class method is a proxy to load_memmap_().

classmethod load_memmap(

prefix: str | Path,

device: device | None = None,

non_blocking: bool = False,

*,

out: TensorDictBase | None = None,

robust_key: bool | None = True,

subpath: NestedKey | None = None,

mode: str = 'r',

num_threads: int = 0,

allow_pickle: bool | None = None,

) → Any

Loads a memory-mapped tensordict from disk.

Parameters:

• prefix (str or Path to folder) – the path to the folder where the saved tensordict should be fetched, or the path to a memmap archive file written through save(..., archive=True) / a ".tdz" prefix (or packed with pack_memmap()). Archives are memory-mapped once and every leaf is exposed as a zero-copy view into the mapping: only the pages of the leaves that are actually accessed are read from disk. Unlike directory-backed tensordicts, in-place writes to the leaves of an archive-loaded tensordict do not propagate to the file.

• device (torch.device or equivalent, optional) – if provided, the data will be asynchronously cast to that device. Supports “meta” device, in which case the data isn’t loaded but a set of empty “meta” tensors are created. This is useful to get a sense of the total model size and structure without actually opening any file.

• non_blocking (bool, optional) – if True, synchronize won’t be called after loading tensors on device. Defaults to False.

• out (TensorDictBase, optional) – optional tensordict where the data should be written.

• robust_key (bool, optional) – if True (default), expects robust key encoding was used when saving and decodes filenames accordingly. If False, uses legacy behavior. If None, uses the default robust behavior.

• subpath (NestedKey or str path, optional) – the location of a nested tensordict to load, as a nested key (e.g. ("module", "0"), with arbitrary nesting allowed as usual) or as a "/"-separated string path (e.g. "module/0"). Only that subtree is loaded. Works both for directories (equivalent to appending the path to prefix) and for archives.

• mode (str, optional) – "r" (default) or "r+". Only relevant when loading an archive: with "r" the archive is mapped copy-on-write and in-place writes to the leaves stay in memory; with "r+" the mapping is shared and in-place writes propagate to the file, like directory-backed tensordicts. "r+" requires uncompressed, aligned tensor payloads (i.e. archives written by tensordict without compression) and is not available for nested-tensor leaves. In-place writes do not update the per-entry CRC-32 stored by the zip format; load_memmap() ignores checksums, but call refresh_archive_checksums() before handing a modified archive to tools that verify them (unzip, unpack_memmap(), …). Directory prefixes are always write-through and ignore this argument.

• num_threads (int, optional) – number of threads used to decompress the leaves of a compressed archive (deflate entries are inflated in parallel, which scales nearly linearly). Without compression, loading is a metadata-only operation and this argument has no effect. Defaults to 0 (sequential).

• allow_pickle (bool, optional) – whether pickled non-tensor fields may be loaded. Pickle can execute arbitrary code, so pass True only for data from a trusted source and False for untrusted data. During the 0.14 compatibility window, omitting this option loads pickle with a FutureWarning; the default will change to False in 0.15. Saves without a pickle sidecar do not require this option.

Examples

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

This method also allows loading nested tensordicts.

Examples

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

A tensordict can also be loaded on “meta” device or, alternatively, as a fake tensor.

Examples

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

state_dict: dict[str, Any],

strict=True,

assign=False,

from_flatten=None,

)

Loads a state_dict into the tensorclass.

Supports both the new format (logical keys with _metadata) and the legacy format (_tensordict/_non_tensordict wrapper keys).

memmap(

prefix: str | None = None,

copy_existing: bool = False,

*,

num_threads: int = 0,

return_early: bool = False,

share_non_tensor: bool = False,

existsok: bool = True,

robust_key: bool | None = True,

archive: bool | None = None,

compression: str | int | None = None,

) → Any

Writes all tensors onto a corresponding memory-mapped Tensor in a new tensordict.

Parameters:

• prefix (str) – directory prefix where the memory-mapped tensors will be stored. The directory tree structure will mimic the tensordict’s. If prefix ends with ".tdz" (or archive=True is passed), a single-file archive is written instead of a directory: a standard zip file whose entries replicate the memmap directory layout. See archive below.

• copy_existing (bool) – If False (default), an exception will be raised if an entry in the tensordict is already a tensor stored on disk with an associated file, but is not saved in the correct location according to prefix. If True, any existing Tensor will be copied to the new location.

Keyword Arguments:

• num_threads (int, optional) – the number of threads used to write the memmap tensors. Defaults to 0.

• return_early (bool, optional) – if True and num_threads>0, the method will return a future of the tensordict.

• share_non_tensor (bool, optional) – if True, the non-tensor data will be shared between the processes and writing operation (such as inplace update or set) on any of the workers within a single node will update the value on all other workers. If the number of non_tensor leaves is high (e.g., sharing large stacks of non-tensor data) this may result in OOM or similar errors. Defaults to False.

• existsok (bool, optional) – if False, an exception will be raised if a tensor already exists in the same path. Defaults to True.

• robust_key (bool, optional) – if True (default), uses robust key encoding that safely handles keys with path separators and special characters. If False, uses legacy behavior (keys used as-is). If None, uses the default robust behavior.

• archive (bool, optional) – if True, prefix designates a single file rather than a directory and the tensordict is written as a memmap archive: a zip file mirroring the memmap directory tree, with tensor payloads stored uncompressed and aligned so that load_memmap() can memory-map the file and expose every leaf as a zero-copy view. If None (default), archive mode is enabled when prefix ends with ".tdz". The result of load_memmap() on an archive behaves like the result of from_consolidated(): all leaves are views into a single storage, and in-place writes do not propagate to the file. Archives and memmap directories are mutually convertible with pack_memmap() / unpack_memmap() (or any zip tool). Note that archives are written sequentially (single data pass) and num_threads has no effect on them.

• compression (str or int, optional) – compression for archive entries ("stored", "deflate", "bzip2", "lzma" or a zipfile constant). Defaults to "stored" (uncompressed), which is what enables zero-copy loading. Compressed archives load correctly but leaves are decompressed in memory on access. Only valid in archive mode.

The TensorDict is then locked, meaning that any writing operations that isn’t in-place will throw an exception (eg, rename, set or remove an entry). Once the tensordict is unlocked, the memory-mapped attribute is turned to False, because cross-process identity is not guaranteed anymore.

Returns:

A new tensordict with the tensors stored on disk if return_early=False, otherwise a TensorDictFuture instance.

Note

Serialising in this fashion might be slow with deeply nested tensordicts, so it is not recommended to call this method inside a training loop.

memmap_(

prefix: str | None = None,

copy_existing: bool = False,

*,

num_threads: int = 0,

return_early: bool = False,

share_non_tensor: bool = False,

existsok: bool = True,

robust_key: bool | None = True,

) → Any

Writes all tensors onto a corresponding memory-mapped Tensor, in-place.

Parameters:

• prefix (str) – directory prefix where the memory-mapped tensors will be stored. The directory tree structure will mimic the tensordict’s.

• copy_existing (bool) – If False (default), an exception will be raised if an entry in the tensordict is already a tensor stored on disk with an associated file, but is not saved in the correct location according to prefix. If True, any existing Tensor will be copied to the new location.

Keyword Arguments:

• num_threads (int, optional) – the number of threads used to write the memmap tensors. Defaults to 0.

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

• share_non_tensor (bool, optional) – if True, the non-tensor data will be shared between the processes and writing operation (such as inplace update or set) on any of the workers within a single node will update the value on all other workers. If the number of non-tensor leaves is high (e.g., sharing large stacks of non-tensor data) this may result in OOM or similar errors. Defaults to False.

• existsok (bool, optional) – if False, an exception will be raised if a tensor already exists in the same path. Defaults to True.

• robust_key (bool, optional) – if True (default), uses robust key encoding that safely handles keys with path separators and special characters. If False, uses legacy behavior (keys used as-is). If None, uses the default robust behavior.

The TensorDict is then locked, meaning that any writing operations that isn’t in-place will throw an exception (eg, rename, set or remove an entry). Once the tensordict is unlocked, the memory-mapped attribute is turned to False, because cross-process identity is not guaranteed anymore.

Returns:

self if return_early=False, otherwise a TensorDictFuture instance.

Note

Serialising in this fashion might be slow with deeply nested tensordicts, so it is not recommended to call this method inside a training loop.

memmap_like(

prefix: str | None = None,

copy_existing: bool = False,

*,

existsok: bool = True,

num_threads: int = 0,

return_early: bool = False,

share_non_tensor: bool = False,

robust_key: bool | None = True,

archive: bool | None = None,

) → Any

Creates a contentless Memory-mapped tensordict with the same shapes as the original one.

Parameters:

• prefix (str) – directory prefix where the memory-mapped tensors will be stored. The directory tree structure will mimic the tensordict’s. If prefix ends with ".tdz" (or archive=True is passed), a preallocated single-file archive is created instead. See archive below.

• copy_existing (bool) – If False (default), an exception will be raised if an entry in the tensordict is already a tensor stored on disk with an associated file, but is not saved in the correct location according to prefix. If True, any existing Tensor will be copied to the new location.

Keyword Arguments:

• num_threads (int, optional) – the number of threads used to write the memmap tensors. Defaults to 0.

• return_early (bool, optional) – if True and num_threads>0, the method will return a future of the tensordict.

• share_non_tensor (bool, optional) – if True, the non-tensor data will be shared between the processes and writing operation (such as inplace update or set) on any of the workers within a single node will update the value on all other workers. If the number of non-tensor leaves is high (e.g., sharing large stacks of non-tensor data) this may result in OOM or similar errors. Defaults to False.

• existsok (bool, optional) – if False, an exception will be raised if a tensor already exists in the same path. Defaults to True.

• robust_key (bool, optional) – if True (default), uses robust key encoding that safely handles keys with path separators and special characters. If False, uses legacy behavior (keys used as-is). If None, uses the default robust behavior.

• archive (bool, optional) – if True, prefix designates a single file and a preallocated, zero-filled memmap archive is created and loaded back with load_memmap(prefix, mode="r+"): the returned tensordict writes through to the archive. If None (default), archive mode is enabled when prefix ends with ".tdz". In-place writes leave the zip per-entry checksums stale; call refresh_archive_checksums() before handing the archive to tools that verify them. Nested tensors are not supported in this mode.

The TensorDict is then locked, meaning that any writing operations that isn’t in-place will throw an exception (eg, rename, set or remove an entry). Once the tensordict is unlocked, the memory-mapped attribute is turned to False, because cross-process identity is not guaranteed anymore.

Returns:

A new TensorDict instance with data stored as memory-mapped tensors if return_early=False, otherwise a TensorDictFuture instance.

Note

This is the recommended method to write a set of large buffers on disk, as memmap_() will copy the information, which can be slow for large content.

Examples

>>> 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_(*, allow_pickle: bool | None = None)

Refreshes the content of the memory-mapped tensordict if it has a saved_path.

This method will raise an exception if no path is associated with it.

Parameters:

allow_pickle (bool, optional) – whether pickled non-tensor fields may be loaded. See load_memmap().

property n_nodes: int

Number of nodes in the tree.

property n_sources: int

Number of source points.

property n_spatial_dims: int

Spatial dimensionality.

save(

prefix: str | None = None,

copy_existing: bool = False,

*,

num_threads: int = 0,

return_early: bool = False,

share_non_tensor: bool = False,

robust_key: bool | None = True,

archive: bool | None = None,

compression: str | int | None = None,

) → Any

Saves the tensordict to disk.

This function is a proxy to memmap().

select(

*keys,

inplace: bool = False,

strict: bool = True,

as_tensordict: bool = False,

)

TensorClass-specific select that supports as_tensordict.

set(

key: NestedKey,

value: Any,

inplace: bool = False,

non_blocking: bool = False,

)

Sets a new key-value pair.

Parameters:

• key (str, tuple of str) – name of the key to be set. If tuple of str it is equivalent to chained calls of getattr followed by a final setattr.

• value (Any) – value to be stored in the tensorclass

• inplace (bool, optional) – if True, set will tentatively try to update the value in-place. If False or if the key isn’t present, the value will be simply written at its destination.

Returns:

self

state_dict(

destination=None,

prefix='',

keep_vars=False,

flatten=True,

) → dict[str, Any]

Returns a state_dict with logical keys, matching TensorDictBase conventions.

Tensor fields appear as data keys. Non-tensor fields (strings, ints, etc.) and the tensorclass type are stored in _metadata. This replaces the legacy _tensordict/_non_tensordict wrapper format.

to_tensordict(

*,

retain_none: bool | None = None,

) → TensorDict

Convert the tensorclass into a regular TensorDict.

Makes a copy of all entries. Memmap and shared memory tensors are converted to regular tensors.

Parameters:

retain_none (bool) – if True, the None values will be written in the tensordict. Otherwise they will be discrarded. Default: True.

Returns:

A new TensorDict object containing the same values as the tensorclass.

unbind(dim: int)

Returns a tuple of indexed tensorclass instances unbound along the indicated dimension.

Resulting tensorclass instances will share the storage of the initial tensorclass instance.

class physicsnemo.mesh.spatial.DualInteractionPlan(

near_target_ids: jaxtyping.Int[Tensor, 'n_near'],

near_source_ids: jaxtyping.Int[Tensor, 'n_near'],

far_target_node_ids: jaxtyping.Int[Tensor, 'n_far_nodes'],

far_source_node_ids: jaxtyping.Int[Tensor, 'n_far_nodes'],

nf_target_ids: jaxtyping.Int[Tensor, 'n_nf'],

nf_source_node_ids: jaxtyping.Int[Tensor, 'n_nf'],

fn_target_node_ids: jaxtyping.Int[Tensor, 'n_fn'],

fn_source_ids: jaxtyping.Int[Tensor, 'n_fn'],

fn_broadcast_targets: jaxtyping.Int[Tensor, 'n_fn_bcast'],

fn_broadcast_starts: jaxtyping.Int[Tensor, 'n_fn'],

fn_broadcast_counts: jaxtyping.Int[Tensor, 'n_fn'],

*,

batch_size,

device=None,

names=None,

)

Bases: object

property device: device

Retrieves the device type of tensor class.

dumps(

prefix: str | None = None,

copy_existing: bool = False,

*,

num_threads: int = 0,

return_early: bool = False,

share_non_tensor: bool = False,

robust_key: bool | None = True,

archive: bool | None = None,

compression: str | int | None = None,

) → Any

Saves the tensordict to disk.

This function is a proxy to memmap().

classmethod fields()

Return a tuple describing the fields of this dataclass.

Accepts a dataclass or an instance of one. Tuple elements are of type Field.

from_csv(

*,

auto_batch_size: bool = False,

batch_dims: int | None = None,

device: device | None = None,

batch_size: Size | None = None,

separator: str | None = None,

dtype: dtype | None = None,

**kwargs,

) → Any

Creates a TensorDict from a CSV file.

Requires either pandas or pyarrow to be installed.

Parameters:

path (str or Path) – Path to the CSV file.

Keyword Arguments:

• auto_batch_size (bool, optional) – If True, the batch size will be computed automatically. Defaults to False.

• batch_dims (int, optional) – If auto_batch_size is True, defines how many dimensions the output tensordict should have. Defaults to None.

• device (torch.device, optional) – The device for tensor data. Defaults to None.

• batch_size (torch.Size, optional) – The batch size. Defaults to [num_rows].

• separator (str, optional) – If provided, column names are split on this separator to create nested TensorDicts. Defaults to None.

• dtype (torch.dtype, optional) – If provided, all numeric columns are cast to this dtype. Defaults to None.

• **kwargs – Additional keyword arguments forwarded to the CSV reader (pandas.read_csv or pyarrow.csv.read_csv).

Returns:

A TensorDict representation of the CSV data.

Examples

>>> td = TensorDict.from_csv("data.csv")
>>> td = TensorDict.from_csv("data.csv", separator=".", dtype=torch.float32)

from_json(

*,

auto_batch_size: bool = False,

batch_dims: int | None = None,

device: device | None = None,

batch_size: Size | None = None,

separator: str | None = None,

dtype: dtype | None = None,

lines: bool = False,

**kwargs,

) → Any

Creates a TensorDict from a JSON file.

Supports both standard JSON (array of records) and JSON Lines format. For nested JSON objects, use from_dict() instead.

Requires pandas for best results. Falls back to stdlib json for simple cases.

Parameters:

path (str or Path) – Path to the JSON file.

Keyword Arguments:

• auto_batch_size (bool, optional) – If True, the batch size will be computed automatically. Defaults to False.

• batch_dims (int, optional) – If auto_batch_size is True, defines how many dimensions the output tensordict should have. Defaults to None.

• device (torch.device, optional) – The device for tensor data. Defaults to None.

• batch_size (torch.Size, optional) – The batch size. Defaults to [num_rows].

• separator (str, optional) – If provided, column names are split on this separator to create nested TensorDicts. Defaults to None.

• dtype (torch.dtype, optional) – If provided, all numeric columns are cast to this dtype. Defaults to None.

• lines (bool, optional) – If True, reads the file as JSON Lines (one JSON object per line). Defaults to False.

• **kwargs – Additional keyword arguments forwarded to the JSON reader.

Returns:

A TensorDict representation of the JSON data.

Examples

>>> td = TensorDict.from_json("data.json")
>>> td = TensorDict.from_json("data.jsonl", lines=True)

from_pandas(

*,

auto_batch_size: bool = False,

batch_dims: int | None = None,

device: device | None = None,

batch_size: Size | None = None,

separator: str | None = None,

dtype: dtype | None = None,

) → Any

Converts a pandas DataFrame to a TensorDict.

Numeric columns become tensors, string/object columns become NonTensorData.

Parameters:

dataframe (pd.DataFrame) – The pandas DataFrame to convert.

Keyword Arguments:

• auto_batch_size (bool, optional) – If True, the batch size will be computed automatically. Defaults to False.

• batch_dims (int, optional) – If auto_batch_size is True, defines how many dimensions the output tensordict should have. Defaults to None.

• device (torch.device, optional) – The device for tensor data. Defaults to None.

• batch_size (torch.Size, optional) – The batch size. Defaults to [num_rows].

• separator (str, optional) – If provided, column names are split on this separator to create nested TensorDicts. For example, with separator=".", a column "obs.x" becomes td["obs", "x"]. Defaults to None.

• dtype (torch.dtype, optional) – If provided, all numeric columns are cast to this dtype. Defaults to None.

Returns:

A TensorDict representation of the DataFrame.

Examples

>>> import pandas as pd
>>> df = pd.DataFrame({"a": [1, 2, 3], "b": [4.0, 5.0, 6.0]})
>>> td = TensorDict.from_pandas(df)
>>> print(td)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.int64, is_shared=False),
        b: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float64, is_shared=False)},
    batch_size=torch.Size([3]),
    device=None,
    is_shared=False)

from_parquet(

*,

auto_batch_size: bool = False,

batch_dims: int | None = None,

device: device | None = None,

batch_size: Size | None = None,

separator: str | None = None,

dtype: dtype | None = None,

columns: list[str] | None = None,

**kwargs,

) → Any

Creates a TensorDict from a Parquet file.

Requires either pyarrow or pandas to be installed. Prefers pyarrow when available for better performance.

Parameters:

path (str or Path) – Path to the Parquet file.

Keyword Arguments:

• auto_batch_size (bool, optional) – If True, the batch size will be computed automatically. Defaults to False.

• batch_dims (int, optional) – If auto_batch_size is True, defines how many dimensions the output tensordict should have. Defaults to None.

• device (torch.device, optional) – The device for tensor data. Defaults to None.

• batch_size (torch.Size, optional) – The batch size. Defaults to [num_rows].

• separator (str, optional) – If provided, column names are split on this separator to create nested TensorDicts. Defaults to None.

• dtype (torch.dtype, optional) – If provided, all numeric columns are cast to this dtype. Defaults to None.

• columns (list of str, optional) – If provided, only read these columns from the file. Defaults to None (all columns).

• **kwargs – Additional keyword arguments forwarded to the Parquet reader.

Returns:

A TensorDict representation of the Parquet data.

Examples

>>> td = TensorDict.from_parquet("data.parquet")
>>> td = TensorDict.from_parquet("data.parquet", columns=["obs", "reward"])

from_schema(

*,

batch_size: Sequence[int] | Size | None = None,

storage: str | None = None,

device=None,

**kwargs,

) → TensorDictBase

Pre-allocate a zero-filled TensorDict from a schema.

Creates a TensorDictBase whose storage backend is selected by storage. Each entry in schema maps a field name to an (element_shape, dtype) pair; the full stored shape is [*batch_size, *element_shape].

Parameters:

schema – Mapping from field name to (element_shape, dtype). element_shape is the per-element shape (excluding batch_size).

Keyword Arguments:

• batch_size – Overall batch dimensions prepended to every element shape. Defaults to ().

• storage (str or None) –

  Backend selector:

  • None – plain TensorDict with regular tensors.

  • "memmap" – memory-mapped tensors on disk. Pass prefix=<dir> in kwargs.

  • "h5" – HDF5 via PersistentTensorDict. Pass filename=<path> in kwargs.

  • "zarr" – zarr (requires zarr>=3.0) via PersistentTensorDict. Pass filename=<path or store> in kwargs.

  • "shared" – CPU shared-memory tensors.

  • "redis" / "dragonfly" – delegates to TensorDictStore.from_schema().

• device – Device for the resulting tensors (ignored by some backends).

• **kwargs – Backend-specific arguments forwarded to the underlying constructor (e.g. prefix for memmap, filename for h5, host/port for redis).

Returns:

A new TensorDictBase subclass instance with pre-allocated (zero-filled) keys.

Examples

>>> td = TensorDict.from_schema(
...     {"obs": ([84, 84, 3], torch.uint8),
...      "reward": ([], torch.float32)},
...     batch_size=[1000],
... )
>>> td["obs"].shape
torch.Size([1000, 84, 84, 3])

>>> import tempfile
>>> with tempfile.TemporaryDirectory() as d:
...     td_mm = TensorDict.from_schema(
...         {"obs": ([4], torch.float32)},
...         batch_size=[8],
...         storage="memmap",
...         prefix=d,
...     )
...     assert td_mm.is_memmap()

classmethod from_tensordict(

tensordict: TensorDictBase,

non_tensordict: dict | None = None,

safe: bool = True,

) → Any

Tensor class wrapper to instantiate a new tensor class object.

Parameters:

• tensordict (TensorDictBase) – Dictionary of tensor types

• non_tensordict (dict) – Dictionary with non-tensor and nested tensor class objects

• safe (bool) – Whether to raise an error if the tensordict is not a TensorDictBase instance

get(

key: NestedKey,

*args,

**kwargs,

)

Gets the value stored with the input key.

Parameters:

• key (str, tuple of str) – key to be queried. If tuple of str it is equivalent to chained calls of getattr.

• default – default value if the key is not found in the tensorclass.

Returns:

value stored with the input key

classmethod load(

prefix: str | Path,

*args,

**kwargs,

) → Any

Loads a tensordict from disk.

This class method is a proxy to load_memmap().

load_(

prefix: str | Path,

*args,

**kwargs,

)

Loads a tensordict from disk within the current tensordict.

This class method is a proxy to load_memmap_().

classmethod load_memmap(

prefix: str | Path,

device: device | None = None,

non_blocking: bool = False,

*,

out: TensorDictBase | None = None,

robust_key: bool | None = True,

subpath: NestedKey | None = None,

mode: str = 'r',

num_threads: int = 0,

allow_pickle: bool | None = None,

) → Any

Loads a memory-mapped tensordict from disk.

Parameters:

• prefix (str or Path to folder) – the path to the folder where the saved tensordict should be fetched, or the path to a memmap archive file written through save(..., archive=True) / a ".tdz" prefix (or packed with pack_memmap()). Archives are memory-mapped once and every leaf is exposed as a zero-copy view into the mapping: only the pages of the leaves that are actually accessed are read from disk. Unlike directory-backed tensordicts, in-place writes to the leaves of an archive-loaded tensordict do not propagate to the file.

• device (torch.device or equivalent, optional) – if provided, the data will be asynchronously cast to that device. Supports “meta” device, in which case the data isn’t loaded but a set of empty “meta” tensors are created. This is useful to get a sense of the total model size and structure without actually opening any file.

• non_blocking (bool, optional) – if True, synchronize won’t be called after loading tensors on device. Defaults to False.

• out (TensorDictBase, optional) – optional tensordict where the data should be written.

• robust_key (bool, optional) – if True (default), expects robust key encoding was used when saving and decodes filenames accordingly. If False, uses legacy behavior. If None, uses the default robust behavior.

• subpath (NestedKey or str path, optional) – the location of a nested tensordict to load, as a nested key (e.g. ("module", "0"), with arbitrary nesting allowed as usual) or as a "/"-separated string path (e.g. "module/0"). Only that subtree is loaded. Works both for directories (equivalent to appending the path to prefix) and for archives.

• mode (str, optional) – "r" (default) or "r+". Only relevant when loading an archive: with "r" the archive is mapped copy-on-write and in-place writes to the leaves stay in memory; with "r+" the mapping is shared and in-place writes propagate to the file, like directory-backed tensordicts. "r+" requires uncompressed, aligned tensor payloads (i.e. archives written by tensordict without compression) and is not available for nested-tensor leaves. In-place writes do not update the per-entry CRC-32 stored by the zip format; load_memmap() ignores checksums, but call refresh_archive_checksums() before handing a modified archive to tools that verify them (unzip, unpack_memmap(), …). Directory prefixes are always write-through and ignore this argument.

• num_threads (int, optional) – number of threads used to decompress the leaves of a compressed archive (deflate entries are inflated in parallel, which scales nearly linearly). Without compression, loading is a metadata-only operation and this argument has no effect. Defaults to 0 (sequential).

• allow_pickle (bool, optional) – whether pickled non-tensor fields may be loaded. Pickle can execute arbitrary code, so pass True only for data from a trusted source and False for untrusted data. During the 0.14 compatibility window, omitting this option loads pickle with a FutureWarning; the default will change to False in 0.15. Saves without a pickle sidecar do not require this option.

Examples

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

This method also allows loading nested tensordicts.

Examples

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

A tensordict can also be loaded on “meta” device or, alternatively, as a fake tensor.

Examples

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

state_dict: dict[str, Any],

strict=True,

assign=False,

from_flatten=None,

)

Loads a state_dict into the tensorclass.

Supports both the new format (logical keys with _metadata) and the legacy format (_tensordict/_non_tensordict wrapper keys).

memmap(

prefix: str | None = None,

copy_existing: bool = False,

*,

num_threads: int = 0,

return_early: bool = False,

share_non_tensor: bool = False,

existsok: bool = True,

robust_key: bool | None = True,

archive: bool | None = None,

compression: str | int | None = None,

) → Any

Writes all tensors onto a corresponding memory-mapped Tensor in a new tensordict.

Parameters:

• prefix (str) – directory prefix where the memory-mapped tensors will be stored. The directory tree structure will mimic the tensordict’s. If prefix ends with ".tdz" (or archive=True is passed), a single-file archive is written instead of a directory: a standard zip file whose entries replicate the memmap directory layout. See archive below.

• copy_existing (bool) – If False (default), an exception will be raised if an entry in the tensordict is already a tensor stored on disk with an associated file, but is not saved in the correct location according to prefix. If True, any existing Tensor will be copied to the new location.

Keyword Arguments:

• num_threads (int, optional) – the number of threads used to write the memmap tensors. Defaults to 0.

• return_early (bool, optional) – if True and num_threads>0, the method will return a future of the tensordict.

• share_non_tensor (bool, optional) – if True, the non-tensor data will be shared between the processes and writing operation (such as inplace update or set) on any of the workers within a single node will update the value on all other workers. If the number of non_tensor leaves is high (e.g., sharing large stacks of non-tensor data) this may result in OOM or similar errors. Defaults to False.

• existsok (bool, optional) – if False, an exception will be raised if a tensor already exists in the same path. Defaults to True.

• robust_key (bool, optional) – if True (default), uses robust key encoding that safely handles keys with path separators and special characters. If False, uses legacy behavior (keys used as-is). If None, uses the default robust behavior.

• archive (bool, optional) – if True, prefix designates a single file rather than a directory and the tensordict is written as a memmap archive: a zip file mirroring the memmap directory tree, with tensor payloads stored uncompressed and aligned so that load_memmap() can memory-map the file and expose every leaf as a zero-copy view. If None (default), archive mode is enabled when prefix ends with ".tdz". The result of load_memmap() on an archive behaves like the result of from_consolidated(): all leaves are views into a single storage, and in-place writes do not propagate to the file. Archives and memmap directories are mutually convertible with pack_memmap() / unpack_memmap() (or any zip tool). Note that archives are written sequentially (single data pass) and num_threads has no effect on them.

• compression (str or int, optional) – compression for archive entries ("stored", "deflate", "bzip2", "lzma" or a zipfile constant). Defaults to "stored" (uncompressed), which is what enables zero-copy loading. Compressed archives load correctly but leaves are decompressed in memory on access. Only valid in archive mode.

The TensorDict is then locked, meaning that any writing operations that isn’t in-place will throw an exception (eg, rename, set or remove an entry). Once the tensordict is unlocked, the memory-mapped attribute is turned to False, because cross-process identity is not guaranteed anymore.

Returns:

A new tensordict with the tensors stored on disk if return_early=False, otherwise a TensorDictFuture instance.

Note

Serialising in this fashion might be slow with deeply nested tensordicts, so it is not recommended to call this method inside a training loop.

memmap_(

prefix: str | None = None,

copy_existing: bool = False,

*,

num_threads: int = 0,

return_early: bool = False,

share_non_tensor: bool = False,

existsok: bool = True,

robust_key: bool | None = True,

) → Any

Writes all tensors onto a corresponding memory-mapped Tensor, in-place.

Parameters:

• prefix (str) – directory prefix where the memory-mapped tensors will be stored. The directory tree structure will mimic the tensordict’s.

• copy_existing (bool) – If False (default), an exception will be raised if an entry in the tensordict is already a tensor stored on disk with an associated file, but is not saved in the correct location according to prefix. If True, any existing Tensor will be copied to the new location.

Keyword Arguments:

• num_threads (int, optional) – the number of threads used to write the memmap tensors. Defaults to 0.

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

• share_non_tensor (bool, optional) – if True, the non-tensor data will be shared between the processes and writing operation (such as inplace update or set) on any of the workers within a single node will update the value on all other workers. If the number of non-tensor leaves is high (e.g., sharing large stacks of non-tensor data) this may result in OOM or similar errors. Defaults to False.

• existsok (bool, optional) – if False, an exception will be raised if a tensor already exists in the same path. Defaults to True.

• robust_key (bool, optional) – if True (default), uses robust key encoding that safely handles keys with path separators and special characters. If False, uses legacy behavior (keys used as-is). If None, uses the default robust behavior.

The TensorDict is then locked, meaning that any writing operations that isn’t in-place will throw an exception (eg, rename, set or remove an entry). Once the tensordict is unlocked, the memory-mapped attribute is turned to False, because cross-process identity is not guaranteed anymore.

Returns:

self if return_early=False, otherwise a TensorDictFuture instance.

Note

Serialising in this fashion might be slow with deeply nested tensordicts, so it is not recommended to call this method inside a training loop.

memmap_like(

prefix: str | None = None,

copy_existing: bool = False,

*,

existsok: bool = True,

num_threads: int = 0,

return_early: bool = False,

share_non_tensor: bool = False,

robust_key: bool | None = True,

archive: bool | None = None,

) → Any

Creates a contentless Memory-mapped tensordict with the same shapes as the original one.

Parameters:

• prefix (str) – directory prefix where the memory-mapped tensors will be stored. The directory tree structure will mimic the tensordict’s. If prefix ends with ".tdz" (or archive=True is passed), a preallocated single-file archive is created instead. See archive below.

• copy_existing (bool) – If False (default), an exception will be raised if an entry in the tensordict is already a tensor stored on disk with an associated file, but is not saved in the correct location according to prefix. If True, any existing Tensor will be copied to the new location.

Keyword Arguments:

• num_threads (int, optional) – the number of threads used to write the memmap tensors. Defaults to 0.

• return_early (bool, optional) – if True and num_threads>0, the method will return a future of the tensordict.

• share_non_tensor (bool, optional) – if True, the non-tensor data will be shared between the processes and writing operation (such as inplace update or set) on any of the workers within a single node will update the value on all other workers. If the number of non-tensor leaves is high (e.g., sharing large stacks of non-tensor data) this may result in OOM or similar errors. Defaults to False.

• existsok (bool, optional) – if False, an exception will be raised if a tensor already exists in the same path. Defaults to True.

• robust_key (bool, optional) – if True (default), uses robust key encoding that safely handles keys with path separators and special characters. If False, uses legacy behavior (keys used as-is). If None, uses the default robust behavior.

• archive (bool, optional) – if True, prefix designates a single file and a preallocated, zero-filled memmap archive is created and loaded back with load_memmap(prefix, mode="r+"): the returned tensordict writes through to the archive. If None (default), archive mode is enabled when prefix ends with ".tdz". In-place writes leave the zip per-entry checksums stale; call refresh_archive_checksums() before handing the archive to tools that verify them. Nested tensors are not supported in this mode.

The TensorDict is then locked, meaning that any writing operations that isn’t in-place will throw an exception (eg, rename, set or remove an entry). Once the tensordict is unlocked, the memory-mapped attribute is turned to False, because cross-process identity is not guaranteed anymore.

Returns:

A new TensorDict instance with data stored as memory-mapped tensors if return_early=False, otherwise a TensorDictFuture instance.

Note

This is the recommended method to write a set of large buffers on disk, as memmap_() will copy the information, which can be slow for large content.

Examples

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

*,

allow_pickle: bool | None = None,

)

Refreshes the content of the memory-mapped tensordict if it has a saved_path.

This method will raise an exception if no path is associated with it.

Parameters:

allow_pickle (bool, optional) – whether pickled non-tensor fields may be loaded. See load_memmap().

property n_far_nodes: int

Number of (far,far) node-to-node pairs (each = one kernel eval).

property n_fn: int

Number of (far,near) target-node-to-source-point pairs.

property n_near: int

Number of (near,near) exact individual interaction pairs.

property n_nf: int

Number of (near,far) target-point-to-source-node pairs.

save(

prefix: str | None = None,

copy_existing: bool = False,

*,

num_threads: int = 0,

return_early: bool = False,

share_non_tensor: bool = False,

robust_key: bool | None = True,

archive: bool | None = None,

compression: str | int | None = None,

) → Any

Saves the tensordict to disk.

This function is a proxy to memmap().

select(

*keys,

inplace: bool = False,

strict: bool = True,

as_tensordict: bool = False,

)

TensorClass-specific select that supports as_tensordict.

set(

key: NestedKey,

value: Any,

inplace: bool = False,

non_blocking: bool = False,

)

Sets a new key-value pair.

Parameters:

• key (str, tuple of str) – name of the key to be set. If tuple of str it is equivalent to chained calls of getattr followed by a final setattr.

• value (Any) – value to be stored in the tensorclass

• inplace (bool, optional) – if True, set will tentatively try to update the value in-place. If False or if the key isn’t present, the value will be simply written at its destination.

Returns:

self

state_dict(

destination=None,

prefix='',

keep_vars=False,

flatten=True,

) → dict[str, Any]

Returns a state_dict with logical keys, matching TensorDictBase conventions.

Tensor fields appear as data keys. Non-tensor fields (strings, ints, etc.) and the tensorclass type are stored in _metadata. This replaces the legacy _tensordict/_non_tensordict wrapper format.

to_tensordict(

*,

retain_none: bool | None = None,

) → TensorDict

Convert the tensorclass into a regular TensorDict.

Makes a copy of all entries. Memmap and shared memory tensors are converted to regular tensors.

Parameters:

retain_none (bool) – if True, the None values will be written in the tensordict. Otherwise they will be discrarded. Default: True.

Returns:

A new TensorDict object containing the same values as the tensorclass.

unbind(dim: int)

Returns a tuple of indexed tensorclass instances unbound along the indicated dimension.

Resulting tensorclass instances will share the storage of the initial tensorclass instance.

validate() → None

Check internal consistency of the interaction plan.

Verifies shape pairing, non-negativity, and fn_broadcast bounds. Raises ValueError on any inconsistency. Intended to be called behind a not torch.compiler.is_compiling() guard so it is zero-cost under torch.compile.

Raises:

ValueError – If any internal consistency check fails.

class physicsnemo.mesh.spatial.SourceAggregates(

node_centroid: jaxtyping.Float[Tensor, 'n_nodes n_dims'],

node_source_data: tensordict._td.TensorDict | None,

*,

batch_size,

device=None,

names=None,

)

Bases: object

property device: device

Retrieves the device type of tensor class.

dumps(

prefix: str | None = None,

copy_existing: bool = False,

*,

num_threads: int = 0,

return_early: bool = False,

share_non_tensor: bool = False,

robust_key: bool | None = True,

archive: bool | None = None,

compression: str | int | None = None,

) → Any

Saves the tensordict to disk.

This function is a proxy to memmap().

classmethod fields()

Return a tuple describing the fields of this dataclass.

Accepts a dataclass or an instance of one. Tuple elements are of type Field.

from_csv(

*,

auto_batch_size: bool = False,

batch_dims: int | None = None,

device: device | None = None,

batch_size: Size | None = None,

separator: str | None = None,

dtype: dtype | None = None,

**kwargs,

) → Any

Creates a TensorDict from a CSV file.

Requires either pandas or pyarrow to be installed.

Parameters:

path (str or Path) – Path to the CSV file.

Keyword Arguments:

• auto_batch_size (bool, optional) – If True, the batch size will be computed automatically. Defaults to False.

• batch_dims (int, optional) – If auto_batch_size is True, defines how many dimensions the output tensordict should have. Defaults to None.

• device (torch.device, optional) – The device for tensor data. Defaults to None.

• batch_size (torch.Size, optional) – The batch size. Defaults to [num_rows].

• separator (str, optional) – If provided, column names are split on this separator to create nested TensorDicts. Defaults to None.

• dtype (torch.dtype, optional) – If provided, all numeric columns are cast to this dtype. Defaults to None.

• **kwargs – Additional keyword arguments forwarded to the CSV reader (pandas.read_csv or pyarrow.csv.read_csv).

Returns:

A TensorDict representation of the CSV data.

Examples

>>> td = TensorDict.from_csv("data.csv")
>>> td = TensorDict.from_csv("data.csv", separator=".", dtype=torch.float32)

from_json(

*,

auto_batch_size: bool = False,

batch_dims: int | None = None,

device: device | None = None,

batch_size: Size | None = None,

separator: str | None = None,

dtype: dtype | None = None,

lines: bool = False,

**kwargs,

) → Any

Creates a TensorDict from a JSON file.

Supports both standard JSON (array of records) and JSON Lines format. For nested JSON objects, use from_dict() instead.

Requires pandas for best results. Falls back to stdlib json for simple cases.

Parameters:

path (str or Path) – Path to the JSON file.

Keyword Arguments:

• auto_batch_size (bool, optional) – If True, the batch size will be computed automatically. Defaults to False.

• batch_dims (int, optional) – If auto_batch_size is True, defines how many dimensions the output tensordict should have. Defaults to None.

• device (torch.device, optional) – The device for tensor data. Defaults to None.

• batch_size (torch.Size, optional) – The batch size. Defaults to [num_rows].

• separator (str, optional) – If provided, column names are split on this separator to create nested TensorDicts. Defaults to None.

• dtype (torch.dtype, optional) – If provided, all numeric columns are cast to this dtype. Defaults to None.

• lines (bool, optional) – If True, reads the file as JSON Lines (one JSON object per line). Defaults to False.

• **kwargs – Additional keyword arguments forwarded to the JSON reader.

Returns:

A TensorDict representation of the JSON data.

Examples

>>> td = TensorDict.from_json("data.json")
>>> td = TensorDict.from_json("data.jsonl", lines=True)

from_pandas(

*,

auto_batch_size: bool = False,

batch_dims: int | None = None,

device: device | None = None,

batch_size: Size | None = None,

separator: str | None = None,

dtype: dtype | None = None,

) → Any

Converts a pandas DataFrame to a TensorDict.

Numeric columns become tensors, string/object columns become NonTensorData.

Parameters:

dataframe (pd.DataFrame) – The pandas DataFrame to convert.

Keyword Arguments:

• auto_batch_size (bool, optional) – If True, the batch size will be computed automatically. Defaults to False.

• batch_dims (int, optional) – If auto_batch_size is True, defines how many dimensions the output tensordict should have. Defaults to None.

• device (torch.device, optional) – The device for tensor data. Defaults to None.

• batch_size (torch.Size, optional) – The batch size. Defaults to [num_rows].

• separator (str, optional) – If provided, column names are split on this separator to create nested TensorDicts. For example, with separator=".", a column "obs.x" becomes td["obs", "x"]. Defaults to None.

• dtype (torch.dtype, optional) – If provided, all numeric columns are cast to this dtype. Defaults to None.

Returns:

A TensorDict representation of the DataFrame.

Examples

>>> import pandas as pd
>>> df = pd.DataFrame({"a": [1, 2, 3], "b": [4.0, 5.0, 6.0]})
>>> td = TensorDict.from_pandas(df)
>>> print(td)
TensorDict(
    fields={
        a: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.int64, is_shared=False),
        b: Tensor(shape=torch.Size([3]), device=cpu, dtype=torch.float64, is_shared=False)},
    batch_size=torch.Size([3]),
    device=None,
    is_shared=False)

from_parquet(

*,

auto_batch_size: bool = False,

batch_dims: int | None = None,

device: device | None = None,

batch_size: Size | None = None,

separator: str | None = None,

dtype: dtype | None = None,

columns: list[str] | None = None,

**kwargs,

) → Any

Creates a TensorDict from a Parquet file.

Requires either pyarrow or pandas to be installed. Prefers pyarrow when available for better performance.

Parameters:

path (str or Path) – Path to the Parquet file.

Keyword Arguments:

• auto_batch_size (bool, optional) – If True, the batch size will be computed automatically. Defaults to False.

• batch_dims (int, optional) – If auto_batch_size is True, defines how many dimensions the output tensordict should have. Defaults to None.

• device (torch.device, optional) – The device for tensor data. Defaults to None.

• batch_size (torch.Size, optional) – The batch size. Defaults to [num_rows].

• separator (str, optional) – If provided, column names are split on this separator to create nested TensorDicts. Defaults to None.

• dtype (torch.dtype, optional) – If provided, all numeric columns are cast to this dtype. Defaults to None.

• columns (list of str, optional) – If provided, only read these columns from the file. Defaults to None (all columns).

• **kwargs – Additional keyword arguments forwarded to the Parquet reader.

Returns:

A TensorDict representation of the Parquet data.

Examples

>>> td = TensorDict.from_parquet("data.parquet")
>>> td = TensorDict.from_parquet("data.parquet", columns=["obs", "reward"])

from_schema(

*,

batch_size: Sequence[int] | Size | None = None,

storage: str | None = None,

device=None,

**kwargs,

) → TensorDictBase

Pre-allocate a zero-filled TensorDict from a schema.

Creates a TensorDictBase whose storage backend is selected by storage. Each entry in schema maps a field name to an (element_shape, dtype) pair; the full stored shape is [*batch_size, *element_shape].

Parameters:

schema – Mapping from field name to (element_shape, dtype). element_shape is the per-element shape (excluding batch_size).

Keyword Arguments:

• batch_size – Overall batch dimensions prepended to every element shape. Defaults to ().

• storage (str or None) –

  Backend selector:

  • None – plain TensorDict with regular tensors.

  • "memmap" – memory-mapped tensors on disk. Pass prefix=<dir> in kwargs.

  • "h5" – HDF5 via PersistentTensorDict. Pass filename=<path> in kwargs.

  • "zarr" – zarr (requires zarr>=3.0) via PersistentTensorDict. Pass filename=<path or store> in kwargs.

  • "shared" – CPU shared-memory tensors.

  • "redis" / "dragonfly" – delegates to TensorDictStore.from_schema().

• device – Device for the resulting tensors (ignored by some backends).

• **kwargs – Backend-specific arguments forwarded to the underlying constructor (e.g. prefix for memmap, filename for h5, host/port for redis).

Returns:

A new TensorDictBase subclass instance with pre-allocated (zero-filled) keys.

Examples

>>> td = TensorDict.from_schema(
...     {"obs": ([84, 84, 3], torch.uint8),
...      "reward": ([], torch.float32)},
...     batch_size=[1000],
... )
>>> td["obs"].shape
torch.Size([1000, 84, 84, 3])

>>> import tempfile
>>> with tempfile.TemporaryDirectory() as d:
...     td_mm = TensorDict.from_schema(
...         {"obs": ([4], torch.float32)},
...         batch_size=[8],
...         storage="memmap",
...         prefix=d,
...     )
...     assert td_mm.is_memmap()

classmethod from_tensordict(

tensordict: TensorDictBase,

non_tensordict: dict | None = None,

safe: bool = True,

) → Any

Tensor class wrapper to instantiate a new tensor class object.

Parameters:

• tensordict (TensorDictBase) – Dictionary of tensor types

• non_tensordict (dict) – Dictionary with non-tensor and nested tensor class objects

• safe (bool) – Whether to raise an error if the tensordict is not a TensorDictBase instance

get(

key: NestedKey,

*args,

**kwargs,

)

Gets the value stored with the input key.

Parameters:

• key (str, tuple of str) – key to be queried. If tuple of str it is equivalent to chained calls of getattr.

• default – default value if the key is not found in the tensorclass.

Returns:

value stored with the input key

classmethod load(

prefix: str | Path,

*args,

**kwargs,

) → Any

Loads a tensordict from disk.

This class method is a proxy to load_memmap().

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

Loads a tensordict from disk within the current tensordict.

This class method is a proxy to load_memmap_().

classmethod load_memmap(

prefix: str | Path,

device: device | None = None,

non_blocking: bool = False,

*,

out: TensorDictBase | None = None,

robust_key: bool | None = True,

subpath: NestedKey | None = None,

mode: str = 'r',

num_threads: int = 0,

allow_pickle: bool | None = None,

) → Any

Loads a memory-mapped tensordict from disk.

Parameters:

• prefix (str or Path to folder) – the path to the folder where the saved tensordict should be fetched, or the path to a memmap archive file written through save(..., archive=True) / a ".tdz" prefix (or packed with pack_memmap()). Archives are memory-mapped once and every leaf is exposed as a zero-copy view into the mapping: only the pages of the leaves that are actually accessed are read from disk. Unlike directory-backed tensordicts, in-place writes to the leaves of an archive-loaded tensordict do not propagate to the file.

• device (torch.device or equivalent, optional) – if provided, the data will be asynchronously cast to that device. Supports “meta” device, in which case the data isn’t loaded but a set of empty “meta” tensors are created. This is useful to get a sense of the total model size and structure without actually opening any file.

• non_blocking (bool, optional) – if True, synchronize won’t be called after loading tensors on device. Defaults to False.

• out (TensorDictBase, optional) – optional tensordict where the data should be written.

• robust_key (bool, optional) – if True (default), expects robust key encoding was used when saving and decodes filenames accordingly. If False, uses legacy behavior. If None, uses the default robust behavior.

• subpath (NestedKey or str path, optional) – the location of a nested tensordict to load, as a nested key (e.g. ("module", "0"), with arbitrary nesting allowed as usual) or as a "/"-separated string path (e.g. "module/0"). Only that subtree is loaded. Works both for directories (equivalent to appending the path to prefix) and for archives.

• mode (str, optional) – "r" (default) or "r+". Only relevant when loading an archive: with "r" the archive is mapped copy-on-write and in-place writes to the leaves stay in memory; with "r+" the mapping is shared and in-place writes propagate to the file, like directory-backed tensordicts. "r+" requires uncompressed, aligned tensor payloads (i.e. archives written by tensordict without compression) and is not available for nested-tensor leaves. In-place writes do not update the per-entry CRC-32 stored by the zip format; load_memmap() ignores checksums, but call refresh_archive_checksums() before handing a modified archive to tools that verify them (unzip, unpack_memmap(), …). Directory prefixes are always write-through and ignore this argument.

• num_threads (int, optional) – number of threads used to decompress the leaves of a compressed archive (deflate entries are inflated in parallel, which scales nearly linearly). Without compression, loading is a metadata-only operation and this argument has no effect. Defaults to 0 (sequential).

• allow_pickle (bool, optional) – whether pickled non-tensor fields may be loaded. Pickle can execute arbitrary code, so pass True only for data from a trusted source and False for untrusted data. During the 0.14 compatibility window, omitting this option loads pickle with a FutureWarning; the default will change to False in 0.15. Saves without a pickle sidecar do not require this option.

Examples

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

This method also allows loading nested tensordicts.

Examples

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

A tensordict can also be loaded on “meta” device or, alternatively, as a fake tensor.

Examples

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

state_dict: dict[str, Any],

strict=True,

assign=False,

from_flatten=None,

)

Loads a state_dict into the tensorclass.

Supports both the new format (logical keys with _metadata) and the legacy format (_tensordict/_non_tensordict wrapper keys).

memmap(

prefix: str | None = None,

copy_existing: bool = False,

*,

num_threads: int = 0,

return_early: bool = False,

share_non_tensor: bool = False,

existsok: bool = True,

robust_key: bool | None = True,

archive: bool | None = None,

compression: str | int | None = None,

) → Any

Writes all tensors onto a corresponding memory-mapped Tensor in a new tensordict.

Parameters:

• prefix (str) – directory prefix where the memory-mapped tensors will be stored. The directory tree structure will mimic the tensordict’s. If prefix ends with ".tdz" (or archive=True is passed), a single-file archive is written instead of a directory: a standard zip file whose entries replicate the memmap directory layout. See archive below.

• copy_existing (bool) – If False (default), an exception will be raised if an entry in the tensordict is already a tensor stored on disk with an associated file, but is not saved in the correct location according to prefix. If True, any existing Tensor will be copied to the new location.

Keyword Arguments:

• num_threads (int, optional) – the number of threads used to write the memmap tensors. Defaults to 0.

• return_early (bool, optional) – if True and num_threads>0, the method will return a future of the tensordict.

• share_non_tensor (bool, optional) – if True, the non-tensor data will be shared between the processes and writing operation (such as inplace update or set) on any of the workers within a single node will update the value on all other workers. If the number of non_tensor leaves is high (e.g., sharing large stacks of non-tensor data) this may result in OOM or similar errors. Defaults to False.

• existsok (bool, optional) – if False, an exception will be raised if a tensor already exists in the same path. Defaults to True.

• robust_key (bool, optional) – if True (default), uses robust key encoding that safely handles keys with path separators and special characters. If False, uses legacy behavior (keys used as-is). If None, uses the default robust behavior.

• archive (bool, optional) – if True, prefix designates a single file rather than a directory and the tensordict is written as a memmap archive: a zip file mirroring the memmap directory tree, with tensor payloads stored uncompressed and aligned so that load_memmap() can memory-map the file and expose every leaf as a zero-copy view. If None (default), archive mode is enabled when prefix ends with ".tdz". The result of load_memmap() on an archive behaves like the result of from_consolidated(): all leaves are views into a single storage, and in-place writes do not propagate to the file. Archives and memmap directories are mutually convertible with pack_memmap() / unpack_memmap() (or any zip tool). Note that archives are written sequentially (single data pass) and num_threads has no effect on them.

• compression (str or int, optional) – compression for archive entries ("stored", "deflate", "bzip2", "lzma" or a zipfile constant). Defaults to "stored" (uncompressed), which is what enables zero-copy loading. Compressed archives load correctly but leaves are decompressed in memory on access. Only valid in archive mode.

The TensorDict is then locked, meaning that any writing operations that isn’t in-place will throw an exception (eg, rename, set or remove an entry). Once the tensordict is unlocked, the memory-mapped attribute is turned to False, because cross-process identity is not guaranteed anymore.

Returns:

A new tensordict with the tensors stored on disk if return_early=False, otherwise a TensorDictFuture instance.

Note

Serialising in this fashion might be slow with deeply nested tensordicts, so it is not recommended to call this method inside a training loop.

memmap_(

prefix: str | None = None,

copy_existing: bool = False,

*,

num_threads: int = 0,

return_early: bool = False,

share_non_tensor: bool = False,

existsok: bool = True,

robust_key: bool | None = True,

) → Any

Writes all tensors onto a corresponding memory-mapped Tensor, in-place.

Parameters:

• prefix (str) – directory prefix where the memory-mapped tensors will be stored. The directory tree structure will mimic the tensordict’s.

• copy_existing (bool) – If False (default), an exception will be raised if an entry in the tensordict is already a tensor stored on disk with an associated file, but is not saved in the correct location according to prefix. If True, any existing Tensor will be copied to the new location.

Keyword Arguments:

• num_threads (int, optional) – the number of threads used to write the memmap tensors. Defaults to 0.

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

• share_non_tensor (bool, optional) – if True, the non-tensor data will be shared between the processes and writing operation (such as inplace update or set) on any of the workers within a single node will update the value on all other workers. If the number of non-tensor leaves is high (e.g., sharing large stacks of non-tensor data) this may result in OOM or similar errors. Defaults to False.

• existsok (bool, optional) – if False, an exception will be raised if a tensor already exists in the same path. Defaults to True.

• robust_key (bool, optional) – if True (default), uses robust key encoding that safely handles keys with path separators and special characters. If False, uses legacy behavior (keys used as-is). If None, uses the default robust behavior.

The TensorDict is then locked, meaning that any writing operations that isn’t in-place will throw an exception (eg, rename, set or remove an entry). Once the tensordict is unlocked, the memory-mapped attribute is turned to False, because cross-process identity is not guaranteed anymore.

Returns:

self if return_early=False, otherwise a TensorDictFuture instance.

Note

Serialising in this fashion might be slow with deeply nested tensordicts, so it is not recommended to call this method inside a training loop.

memmap_like(

prefix: str | None = None,

copy_existing: bool = False,

*,

existsok: bool = True,

num_threads: int = 0,

return_early: bool = False,

share_non_tensor: bool = False,

robust_key: bool | None = True,

archive: bool | None = None,

) → Any

Creates a contentless Memory-mapped tensordict with the same shapes as the original one.

Parameters:

• prefix (str) – directory prefix where the memory-mapped tensors will be stored. The directory tree structure will mimic the tensordict’s. If prefix ends with ".tdz" (or archive=True is passed), a preallocated single-file archive is created instead. See archive below.

• copy_existing (bool) – If False (default), an exception will be raised if an entry in the tensordict is already a tensor stored on disk with an associated file, but is not saved in the correct location according to prefix. If True, any existing Tensor will be copied to the new location.

Keyword Arguments:

• num_threads (int, optional) – the number of threads used to write the memmap tensors. Defaults to 0.

• return_early (bool, optional) – if True and num_threads>0, the method will return a future of the tensordict.

• share_non_tensor (bool, optional) – if True, the non-tensor data will be shared between the processes and writing operation (such as inplace update or set) on any of the workers within a single node will update the value on all other workers. If the number of non-tensor leaves is high (e.g., sharing large stacks of non-tensor data) this may result in OOM or similar errors. Defaults to False.

• existsok (bool, optional) – if False, an exception will be raised if a tensor already exists in the same path. Defaults to True.

• robust_key (bool, optional) – if True (default), uses robust key encoding that safely handles keys with path separators and special characters. If False, uses legacy behavior (keys used as-is). If None, uses the default robust behavior.

• archive (bool, optional) – if True, prefix designates a single file and a preallocated, zero-filled memmap archive is created and loaded back with load_memmap(prefix, mode="r+"): the returned tensordict writes through to the archive. If None (default), archive mode is enabled when prefix ends with ".tdz". In-place writes leave the zip per-entry checksums stale; call refresh_archive_checksums() before handing the archive to tools that verify them. Nested tensors are not supported in this mode.

The TensorDict is then locked, meaning that any writing operations that isn’t in-place will throw an exception (eg, rename, set or remove an entry). Once the tensordict is unlocked, the memory-mapped attribute is turned to False, because cross-process identity is not guaranteed anymore.

Returns:

A new TensorDict instance with data stored as memory-mapped tensors if return_early=False, otherwise a TensorDictFuture instance.

Note

This is the recommended method to write a set of large buffers on disk, as memmap_() will copy the information, which can be slow for large content.

Examples

>>> 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_(*, allow_pickle: bool | None = None)

Refreshes the content of the memory-mapped tensordict if it has a saved_path.

This method will raise an exception if no path is associated with it.

Parameters:

allow_pickle (bool, optional) – whether pickled non-tensor fields may be loaded. See load_memmap().

node_centroid: Float[Tensor, 'n_nodes n_dims']

Area-weighted centroid per node.

node_source_data: TensorDict | None

Area-weighted average source features per node, or None if no per-source features. Has batch_size=(n_nodes,).

save(

prefix: str | None = None,

copy_existing: bool = False,

*,

num_threads: int = 0,

return_early: bool = False,

share_non_tensor: bool = False,

robust_key: bool | None = True,

archive: bool | None = None,

compression: str | int | None = None,

) → Any

Saves the tensordict to disk.

This function is a proxy to memmap().

select(

*keys,

inplace: bool = False,

strict: bool = True,

as_tensordict: bool = False,

)

TensorClass-specific select that supports as_tensordict.

set(

key: NestedKey,

value: Any,

inplace: bool = False,

non_blocking: bool = False,

)

Sets a new key-value pair.

Parameters:

• key (str, tuple of str) – name of the key to be set. If tuple of str it is equivalent to chained calls of getattr followed by a final setattr.

• value (Any) – value to be stored in the tensorclass

• inplace (bool, optional) – if True, set will tentatively try to update the value in-place. If False or if the key isn’t present, the value will be simply written at its destination.

Returns:

self

state_dict(

destination=None,

prefix='',

keep_vars=False,

flatten=True,

) → dict[str, Any]

Returns a state_dict with logical keys, matching TensorDictBase conventions.

Tensor fields appear as data keys. Non-tensor fields (strings, ints, etc.) and the tensorclass type are stored in _metadata. This replaces the legacy _tensordict/_non_tensordict wrapper format.

to_tensordict(

*,

retain_none: bool | None = None,

) → TensorDict

Convert the tensorclass into a regular TensorDict.

Makes a copy of all entries. Memmap and shared memory tensors are converted to regular tensors.

Parameters:

retain_none (bool) – if True, the None values will be written in the tensordict. Otherwise they will be discrarded. Default: True.

Returns:

A new TensorDict object containing the same values as the tensorclass.

unbind(dim: int)

Returns a tuple of indexed tensorclass instances unbound along the indicated dimension.

Resulting tensorclass instances will share the storage of the initial tensorclass instance.

physicsnemo.mesh.spatial.signed_distance_field(

mesh: Mesh,

query_points: Float[Tensor, '... 3'],

max_dist: float | None = None,

use_sign_winding_number: bool = False,

) → tuple[Float[Tensor, '...'], Float[Tensor, '... 3'], Int[Tensor, '...']]

Compute the signed distance to a triangle surface mesh.

Returns the signed distance, the closest surface point, and the nearest face index for each query. Delegates to the Warp-backed physicsnemo.nn.functional.signed_distance_field() op, which runs on CPU and CUDA.

Parameters:

• mesh (Mesh) – Triangle surface mesh embedded in 3D: mesh.points has shape (n_vertices, 3) and mesh.cells has shape (n_faces, 3).

• query_points (torch.Tensor) – Query points, shape (..., 3).

• max_dist (float or None, optional) – Maximum search radius for the nearest-triangle query. None (default) searches without bound, so the true nearest triangle is always found; a finite value restricts the search to a band and reports queries farther than it as NaN (both sdf and hit_points) with a hit face of -1.

• use_sign_winding_number (bool, optional) – If True, sign via the generalized winding number (wp.mesh_query_point_sign_winding_number), robust for non-watertight meshes. If False (default), sign via the angle-weighted pseudo-normal of the closest mesh feature (wp.mesh_query_point_sign_normal), which stays correct at sharp/non-convex edges where a single face normal would flip the sign. The mesh should be watertight for reliable signs in the False case.

Returns:

(sdf, hit_points, hit_faces): signed distance per query (shape query_points.shape[:-1]; negative inside, positive outside), the closest point on the mesh per query (shape query_points.shape), and the index into mesh.cells of the nearest face per query (int64, shape query_points.shape[:-1]). Queries beyond a finite max_dist return NaN for the distance and hit point and -1 for the hit face.

Return type:

tuple[torch.Tensor, torch.Tensor, torch.Tensor]

Raises:

ValueError – If mesh is not a triangle surface in 3D (n_spatial_dims == 3 and n_manifold_dims == 2), if query_points does not have a trailing dimension of size 3, if the mesh has no faces (there is no surface to measure distance to), if mesh and query_points are on different devices, or if max_dist is negative.

Notes

A finite max_dist is an opt-in optimization/narrow-band mode: it prunes the search to the given radius and marks out-of-band queries as NaN so a far query is never silently reported as on-surface (sdf == 0). The unbounded default never produces NaN for a non-empty mesh.

Distances are computed in float32 (inputs are cast as needed) and results are cast back to the query_points dtype.

(Near-)degenerate faces – repeated vertices or collinear points, which Warp’s mesh query would otherwise skip – are repaired into equivalent thin-but-valid triangles over the same longest edge before the query, so a query nearest to such a face reports the distance to its segment (to within the repair offset) rather than the distance to some farther valid face. See _repair_degenerate_faces().