PhysicsNeMo-Mesh: inherit directly from TensorClass - #1868
Conversation
b8f893b to
e4fb05d
Compare
e4fb05d to
8a01a3c
Compare
|
Too many files changed for review. ( Bypass the limit by tagging |
Both classes carried an identical 28-line `_load_memmap` override whose only difference was the stock function it captured. Move it to a single `_serialization.install_legacy_memmap_reader`, and take `**kwargs` instead of hard-coding tensordict's keyword-only `robust_key`, so a signature change upstream doesn't break the override. Drops the `isinstance(out, cls)` branch, which was unreachable: tensordict only reaches `_load_memmap` through a dispatch path that passes neither `device` nor `out` (instrumenting the mesh suite records 11 calls, all with both set to None). Also corrects the comment above `_domain_from_tensordict`. The nested TensorDict-to-Mesh coercion it labels as legacy compatibility is load-critical for the *current* format too: the memmap format records nested containers as plain TensorDicts, so without it `DomainMesh.load` of a file written today fails with "`interior` must be a Mesh, got TensorDict".
`Mesh.load(path, device="cuda")` silently returned CPU tensors for any file
written before the TensorClass switch. tensordict resolves the writing class
from the directory's meta.json, and when that class isn't the one it is
loading through, its fallback dispatch forwards only the prefix -- `device`
and `out` never reach the container.
Override the public `load` / `load_memmap` entry points, where those arguments
are still intact, so both layouts honor `device=`. The legacy payload is read
natively and moved as a whole afterwards: tensordict dispatches each nested
container through `_load_memmap` without a device, so a DomainMesh would
otherwise end up with its interior on a different device than its global data.
Two side effects of routing through the requested class rather than the one
named in meta.json:
- `Mesh.load` of a `.pdmsh` (and vice versa) now raises a clear key-mismatch
ValueError instead of silently returning the other container type.
- Passing a container as `out=` no longer fails on an internal attribute
("Cannot set attribute _memmap_prefix"); its storage is filled instead.
The `.pmsh` / `.pdmsh` golden tests parametrize over a decorator-era fixture and a current one, but nothing asserted that the legacy fixture is actually in the legacy layout. Regenerating it would leave two current-format fixtures, both tests would still pass, and the backward-read coverage would be gone without a trace. Assert the payload nesting each fixture is supposed to have. Add a device test over both layouts, which fails on the legacy parameter before the preceding commit. Also share `_serialization_manifest` (copied verbatim between the two modules) as a conftest fixture, and settle on one policy for a missing regeneration helper: fail loudly. The `.pmsh` module skipped instead, which would quietly drop the on-disk-format coverage that these modules exist to provide.
Under `@tensorclass`, `DomainMesh[...]` raised TypeError. Inheriting from TensorClass brings along its configuration subscript, so `DomainMesh["wall"]` now quietly evaluates to `TensorClass_nocast` -- an unrelated class -- instead of failing. `Mesh` already claims that syntax for dimension specialization via its own metaclass; `DomainMesh` has no such parametrization, so restore the error.
The PR description noted it, but the CHANGELOG is what users read. Files written by this release load as plain TensorDict objects in 2.1.x and earlier -- silently, with no error, so the failure surfaces later as a missing attribute. Worth stating for anyone converting datasets to `.pdmsh`.
The `load` override reads as a redundant proxy to `load_memmap`, and deleting it silently reverts the device fix -- TensorClass's metaclass resolves class attributes against the underlying TensorDict, so the stock `load` never sees an overridden `load_memmap`. Say so where someone would go to delete it. Also inline the single-use directory-name constant and trim a docstring.
Both containers guarded `_from_tensordict` against being handed an already-reconstructed instance, which happened when the legacy loader returned one and TensorClass's wrapper re-wrapped it. Routing legacy reads through the public entry points removed that double-wrap, so the guards no longer fire: instrumenting the mesh suite records 224 Mesh and 9 DomainMesh reconstructions with zero guard hits, and the suite passes with both removed. For Mesh that retires the `_from_tensordict` patch entirely. What remains for DomainMesh is the nested-container coercion, which is load-critical. Also make the two sibling metaclass names parallel.
|
/ok to test bbbfc78 |
|
/ok to test 57d3e2d |
|
/ok to test 1744c40 |
mehdiataei
left a comment
mehdiataei
left a comment
There was a problem hiding this comment.
I found a few compatibility issues in the inheritance and serialization paths.
|
Overall, I think the direction makes sense, and the normal The biggest concern is serialization. A For example, a user might define a cached training sample like this: class TrainingSample(TensorClass):
mesh: Mesh
target: torch.Tensor
sample = TrainingSample(
mesh=mesh,
target=torch.tensor(1.0),
)
sample.save(path)
loaded = TrainingSample.load(path)
type(loaded)
# TrainingSample
type(loaded.mesh)
# TensorDict, expected MeshThis round trip preserved the nested Subclass behavior also feels a bit uneven. Existing subclasses lose constructor defaults, some round trips lose the concrete subtype, and dtype conversions can skip or drop fields added by subclasses. Given that this started as a low-priority cleanup, I don’t think these compatibility edges are quite safe enough to merge yet. My suggestion would be to move the compatibility behavior into a shared inherited base or mixin instead of patching methods onto the finished classes. It should preserve the concrete type in the saved metadata, reconstruct nested objects using that recorded type, and apply automatically to subclasses. I’d also add round-trip tests for nested If nested serialization is intentionally outside the supported API, that limitation should be documented clearly. If solving these cases cleanly makes the change much larger, keeping the decorator for now may be the safer option. |
|
Thanks for the review! Personally, the subclass/inheritance changes don't worry me (as inheritance on a decorator-@Tensorclass was always a bit borked), but I do think that the specific example (TrainingExamples in the public comment) is one we should definitely fix. Looks like this is fixed by upstream pytorch/tensordict#1722, which is merged but not yet released. Let's pause on this PR until that goes live. |
|
/ok to test f72d811 |
|
Paused until pytorch/tensordict#1756 merges, at which point we'll use nightly. |
|
/ok to test 71784ea |
|
/ok to test 641471b |
|
/ok to test bbebe08 |
|
/ok to test 0379218 |
2 participants
PhysicsNeMo Pull Request
Description
Replace the
@tensorclass-decoratedMeshandDomainMeshcontainers withdirect
TensorClassinheritance.Meshretains its tensor-only/shadow behaviorand
Mesh[m, s]runtime specialization syntax, whileDomainMeshremainsunsubscriptable.
TensorDict 0.14 writes directly inherited tensorclasses with their type
discriminator and
_tensordict/payload, so nestedMeshandDomainMeshobjects preserve their concrete types. Committed decorator-era
.pmshand.pdmshfixtures continue to reconstruct exact mesh types, including nestedinterior and boundary meshes. Current writer layouts are pinned by committed
JSON manifests.
The compatibility adapters preserve the inherited public loader signatures and
documentation, support preallocated
out=storage, reject adevice=thatconflicts with the output tensors instead of producing a mixed-device result,
and forward TensorDict's explicit
allow_picklesafety policy for non-tensorfields.
Field defaults remain available to generated constructors for fieldless
subclasses. Serialization, functional mesh updates, and floating dtype
conversion preserve concrete subclasses and their additional tensor fields.
Closes #1851.
Tests on Python 3.14 with TensorDict 0.14.0:
Checklist
Dependencies
Satisfied by the stable
tensordict[zarr]>=0.14.0dependency currently onmain. TensorDict 0.14 contains the typed direct-inheritance serializationsupport from pytorch/tensordict#1722.
Review Process
Ready for review.