Skip to content

Mesh: preserve empty tensors in memmap serialization - #1882

Draft
peterdsharpe wants to merge 1 commit into
NVIDIA:mainfrom
peterdsharpe:pr/mesh-memmap-empty-tensors
Draft

Mesh: preserve empty tensors in memmap serialization#1882
peterdsharpe wants to merge 1 commit into
NVIDIA:mainfrom
peterdsharpe:pr/mesh-memmap-empty-tensors

Conversation

@peterdsharpe

@peterdsharpe peterdsharpe commented Jul 28, 2026

Copy link
Copy Markdown
Collaborator

PhysicsNeMo Pull Request

Description

TensorDict records zero-element tensors in memmap metadata but writes no backing file, so loading silently drops them. This can change empty connectivity from (0, 3) to the (0, 1) point-cloud sentinel, remove empty data fields, or crash while restoring empty adjacency and spatial structures. Loading a cached Adjacency in a fresh process can also fail because its tensorclass has not been registered.

This PR:

  • Rebuilds missing zero-element tensors from existing metadata before tensorclass __post_init__ can replace or dereference them.
  • Applies the loader to Mesh, DomainMesh, Adjacency, BVH, ClusterTree, and interaction-plan tensorclasses.
  • Registers Adjacency when Mesh is imported so fresh-process cache loading succeeds.
  • Leaves writing and the on-disk format unchanged; files remain byte-identical and backward compatible.

Empty jagged tensors remain a documented, strict-xfail limitation because their shapes cannot be reconstructed from meta.json alone.

Verification

  • pytest test/mesh: 2357 passed, 704 skipped, 1 xfailed.
  • Every new regression test fails without its corresponding fix.
  • Saved-file SHA-256 digests are unchanged for empty Mesh, DomainMesh, and BVH cases.
  • ruff check, ruff format --check, and docstring coverage checks pass.

Checklist

  • I am familiar with the Contributing Guidelines.
  • New or existing tests cover these changes.
  • The documentation is up to date with these changes.
  • The CHANGELOG.md is up to date with these changes.
  • An issue is linked to this pull request.
  • Model implementation standards are not applicable; no model code is changed.

Dependencies

None.

Review Process

All PRs are reviewed by the PhysicsNeMo team before merging.

Depending on which files are changed, GitHub may automatically assign a maintainer for review.

We are also testing AI-based code review tools (e.g., Greptile), which may add automated comments with a confidence score. This score reflects the AI's assessment of merge readiness and is not a qualitative judgment of the work or an indication that the PR will be accepted or rejected.

AI-generated feedback should be reviewed critically for usefulness. You are not required to respond to every AI comment, but they are intended to help both authors and reviewers. Please react to Greptile comments with 👍 or 👎 to provide feedback on their accuracy.

Two independent Mesh.save/load (memmap) defects.

TensorDict records a zero-element tensor's shape and dtype in meta.json but
writes no backing file, and its loader skips keys whose file is absent. Cells
of shape (0, 3) loaded back as (0, 1) with a different dtype, zero-width
point/cell/global data fields disappeared, and a mesh with no points, or with
an Adjacency cached over a point cloud, failed to load at all. Mesh,
DomainMesh, Adjacency, BVH, ClusterTree, DualInteractionPlan and
SourceAggregates now rebuild those tensors from metadata on load.

Separately, and independent of empty tensors, loading a mesh with a populated
topology cache in a process that had not otherwise imported the adjacency API
raised "Could not find name ...Adjacency" -- the offline-preprocess /
dataloader-worker pattern. Importing Mesh now registers Adjacency with
TensorDict.

Both fixes are read-side only; the on-disk format is unchanged.
@copy-pr-bot

copy-pr-bot Bot commented Jul 28, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@github-actions

Copy link
Copy Markdown
Contributor

CODEOWNERS review map

Current for commit 606d1159a1dd. An approval covers every file listed for that owner; one owner is sufficient for shared files.

@peterdsharpe — 9 file(s)
  • physicsnemo/mesh/domain_mesh.py
  • physicsnemo/mesh/mesh.py
  • physicsnemo/mesh/neighbors/_adjacency.py
  • physicsnemo/mesh/spatial/bvh.py
  • physicsnemo/mesh/spatial/cluster_tree.py
  • physicsnemo/mesh/utilities/_serialization.py
  • test/mesh/mesh/test_serialization.py
  • test/mesh/spatial/test_bvh.py
  • test/mesh/spatial/test_cluster_tree.py

No CODEOWNER

  • CHANGELOG.md

Comment /codeowners-info to refresh.

@peterdsharpe

Copy link
Copy Markdown
Collaborator Author

Update: probably going to fix this here; keeping this draft until we see if this merges: pytorch/tensordict#1750

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant