fix: keep data-dependent tensor sizes symbolic in exports by lgulich · Pull Request #48 · nvidia-isaac/leapp

lgulich · 2026-10-26T22:08:12Z

Problem

torch.nonzero, boolean masks (x[x > 0]), masked_select and torch.unique return tensors whose size depends on the input values. Reading such a size during tracing (x.shape[0], x.size(0), x.numel()) returned a plain Python int, so the traced count was baked into the exported graph:

idx = torch.nonzero(x > 0)            # [N, 2], N depends on x
picked = x[idx[:, 0], idx[:, 1]]
y = picked.reshape(idx.shape[0], 1)   # exported as Reshape(..., [2, 1])

compile_graph(validate=True) then fails as soon as two samples produce different counts: Input shape:{3,1}, requested shape:{2,1}. If all samples happen to give the same count, validation passes and the bundle fails at deployment instead. The TorchScript exporter (torch.onnx.export(..., dynamo=False)) keeps the same code dynamic.

We hit this while exporting a ProtoMotions humanoid policy (Metatron) through LEAPP.

Fix

Tensors produced by those ops, and tensors computed from them, answer size queries with TracedInt (leapp/leapp_graph/datatypes/primitives/traced_int.py). It subclasses both int and TracedData and carries the FX proxy of the size/numel call, so LEAPP's generic traversal, unwrapping, proxy extraction and context checks treat it like other traced data:
- integer arithmetic on it (+ - * // %) is recorded, but only while its node is tracing on the current graph. After capture, or for a replaced graph, it behaves as a plain int. Operands from another active node raise the usual mixed-context error;
- it is passed to ops as its proxy, including inside torch.Size, shape tuples, and getitem/setitem keys. narrow bounds use the integer proxy directly;
- copy/deepcopy return the same value, it pickles as a plain int (__reduce_ex__), and YAML dumps it as an integer.
The existing torch.tensor/as_tensor patch records torch.tensor(x.shape) as scalar_tensor + stack. The TorchScript tracer keeps torch.tensor(traced ints) constant but traces those two.
Tensors with static shapes still return plain ints, so existing graphs are unchanged.
The debugging guide gets a short "Data-dependent shapes" section. len(x) can't be made symbolic, because CPython coerces __len__ to an exact int; the guide says to use x.size(0) instead.
Not included: using a TracedInt as a node input or output.

Testing

New tests/functional_tests/test_data_dependent_shapes.py covers reshape with a nonzero count, size/numel with integer arithmetic, boolean-mask counts, torch.tensor(x.shape), and static shapes staying plain ints. Each test exports with two samples of different counts, validates, and checks InferenceManager against eager on an unseen third count. All five fail without this change.
New tests/unit_tests/test_traced_int.py (19 tests) covers the TracedData views and traversal, every arithmetic op and the recorded target, copy/pickle/YAML, behaviour after capture and after graph replacement, stale and cross-context operands, narrow with a scalar bound, and replay of slice and value assignments.
Full suite against develop: 1079 passed (develop: 1055). The same 9 failures and 78 errors occur on develop too; they are Warp tests whose C++ runtime can't build in my environment.

🤖 Generated with Claude Code

frlai

This is great. I previously had wanted to solve this problem but my solution would have been worse than this one.

frlai · 2026-10-28T02:36:08Z

call this traced int. we should add a new folder to datatypes called primitives and put this there.

Moved to datatypes/primitives/traced_int.py as TracedInt, exported from leapp.leapp_graph.datatypes.

frlai · 2026-10-28T02:36:48Z

+import torch
+
+
+class TracedInt(int):


can this also be a subclass of traced_data at the same time? this will make it so leapp sees it as a tracable data and will propagate the traced state.

An isolated prototype confirmed that this works and allows existing structure traversal, eager unwrapping, proxy extraction, and context discovery to recognize traced sizes. That would reduce separate scalar-specific traversal and validation logic.

as a subclass of traced data this would probably also need to define things properties like _tensor and _data that help return the right data type. I can handle enabling this to become a valid input or output of a leapp node.

Done: TracedInt now subclasses int and TracedData (with tensor/data views), and the scalar-specific find/proxy/owner helpers are gone in favour of the generic traversal. Node I/O is left for you, as you suggested.

frlai · 2026-10-28T02:40:02Z

+    other integers records the operation so derived sizes stay symbolic.
+    """
+
+    def __new__(cls, value: int, proxy, context, owner):


copy.copy() and copy.deepcopy() currently fail because reconstruction calls TracedInt.__new__ without its tracing arguments. This can probably be fixed with explicit copy hooks.

another probe when i was checking this behavior revealed that if someone tries to serialize this value for example save it in a yaml, that might fail. the agent suggested using __reduce_ex__ returning (int, (int(self),)), avoiding serialization of the live tracer.

Added __copy__/__deepcopy__ (same value) and __reduce_ex__ returning (int, (value,)). yaml.safe_dump still rejected the subclass, so I also registered an int representer; both YAML dumpers now write a plain integer. Covered in tests/unit_tests/test_traced_int.py.

frlai · 2026-10-28T02:41:56Z

+        return self._record(operator.mod, self, other) if isinstance(other, int) else NotImplemented
+
+    def __neg__(self):
+        return self._record(operator.sub, 0, self)


_record() currently creates proxies unconditionally. Please check the owning context’s current tracing status before recording arithmetic, validate operands from other contexts, and handle proxies belonging to a replaced graph. Arithmetic after capture finishes should execute normally without appending operations to an inactive graph.

Arithmetic now records only while the owning node is tracing on its current graph and otherwise returns a plain int. Operands from another active node raise the usual mixed-context error, and stale operands (finished node or replaced graph) are used as constants.

frlai · 2026-10-28T02:44:04Z

note to self:

Scalar-aware consumers — if adopting TracedData inheritance

audit consumers that assume every TracedData represents a tensor. For example, _narrow_bound() currently emits .item() on a traced bound’s proxy. A traced integer already represents an integer graph value, so that produces an AttributeError during replay. Scalar bounds should use their proxy directly; tensor bounds should retain the existing handling. Conversion factories and dtype helpers also need checking.

frlai · 2026-10-28T02:44:41Z

additionally this feature probably warrants unit tests.

Added tests/unit_tests/test_traced_int.py (19 tests): traversal, every arithmetic op, copy/pickle/YAML, after capture, after graph replacement, stale and cross-context operands, narrow with a scalar bound, and assignment replay.

frlai · 2026-10-28T02:59:08Z

can you target develop instead of main

torch.nonzero, boolean masks, masked_select and torch.unique return tensors whose size depends on input values. Reading such a size (x.shape[0], x.size(0), x.numel()) returned a plain int, so the traced count was baked into the exported graph. A later input with a different count then failed in ONNX Runtime, e.g. "cannot reshape {3} to {2}". The TorchScript exporter keeps these sizes dynamic. Tensors produced by those ops, and tensors computed from them, now answer size queries with TracedInt: an int that carries the FX proxy of the size query. Integer arithmetic on it is recorded, it is passed to ops as its proxy, and torch.tensor(x.shape) is recorded as scalar_tensor/stack so the exporter traces it. Tensors with static shapes still return plain ints, so their graphs are unchanged. len() always returns a plain int in Python, so it cannot stay symbolic; the debugging guide says to use x.size(0) instead. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Lionel Gulich <lgulich@nvidia.com>

Address review feedback on the data-dependent shape support: - Move TracedInt to leapp/leapp_graph/datatypes/primitives and export it from leapp.leapp_graph.datatypes. - TracedInt now subclasses both int and TracedData, so the generic structure traversal, unwrapping, proxy extraction and context validation recognize it. This replaces the scalar-specific find/proxy helpers and the owning-tensor back reference. - Arithmetic records only while the owning node is tracing and its graph is current; otherwise it returns a plain int. Operands from another active node are rejected, and operands from a finished or replaced graph are used as constants. The same applies to traced integers passed to tensor ops, indexing and torch.tensor(). - Consumers that assumed every TracedData is a tensor now handle traced integers: narrow() bounds use the integer proxy directly instead of .item(), an assigned traced integer becomes scalar_tensor(), and torch.tensor()/np.asarray() of a finished traced integer stay plain. - copy.copy/deepcopy return the same value, pickling stores a plain int, and YAML dumps it as an integer. - Add unit tests for TracedInt. Signed-off-by: Lionel Gulich <lgulich@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Drop TracedInt.__format__, which int.__format__ already provides, and the redundant traced-integer check before _graph_sizes(), which performs the same check itself. Signed-off-by: Lionel Gulich <lgulich@nvidia.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

lgulich · 2026-10-28T08:15:18Z

Rebased onto develop and retargeted.

On your audit note: _narrow_bound now uses an integer bound's proxy directly, an assigned TracedInt is recorded as scalar_tensor, the torch.tensor/as_tensor patch no longer sends scalars through the tensor-only path, and np.asarray works via __array__. A TracedInt inside a setitem key now takes the dynamic index path instead of being baked in. Still open: Warp's TracedWpArray._find_single_traced_data would treat a TracedInt argument as an array source. I left that for the node I/O work.

lgulich mentioned this pull request Oct 27, 2026

fix: broadcast one-element sources of static indexed assignments #49

Draft

frlai reviewed Oct 28, 2026

View reviewed changes

lgulich and others added 2 commits October 28, 2026 07:57

lgulich force-pushed the lgulich/data-dependent-shape-tracing branch from f57c3eb to 191be0f Compare October 28, 2026 07:11

lgulich changed the base branch from main to develop October 28, 2026 07:11

		import torch


		class TracedInt(int):

Conversation

lgulich commented Oct 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Problem

Fix

Testing

Uh oh!

frlai left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

frlai commented Oct 28, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

frlai commented Oct 28, 2026

Uh oh!

lgulich commented Oct 28, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

lgulich commented Oct 26, 2026 •

edited

Loading