[BugFix] Fix dynamic-shape export (#1003): handle scalar SymInt in _parse_batch_size and store batch_dims in pytree context by felixmaximilian · Pull Request #1704 · pytorch/tensordict

felixmaximilian · 2026-05-16T19:37:22Z

Problem

torch.export.export() with dynamic_shapes failed for any module that constructs a TensorDict with an input-dependent batch_size. Two independent bugs caused this:

Bug 1 — _parse_batch_size rejects scalar SymInt (_td.py)

When batch_size=x.shape[0] is passed during compilation, batch_size is a torch.SymInt. The is_compiling() path checks isinstance(batch_size, Number), but torch.SymInt does not subclass numbers.Number, so the value falls through to raise ValueError():

ValueError: batch size was not specified when creating the TensorDict instance

Bug 2 — _tensordict_flatten stores batch_size (may contain SymInt) in the pytree context (_pytree.py)

When the exported module returns a TensorDict, torch.export serializes the output pytree spec. The context stored batch_size: torch.Size which may contain SymInt values. torch.export calls as_python_constant() on every context item — SymInt is not a Python constant, raising:

AsPythonConstantNotImplementedError: SymNodeVariable() is not a constant

Fix

tensordict/_td.py: extend the Number check in _parse_batch_size to also accept torch.SymInt:

elif isinstance(batch_size, (Number, torch.SymInt)):
    return torch.Size([batch_size])

tensordict/_pytree.py: store batch_dims: int (the count, always a Python constant) in the pytree context instead of batch_size. In _tensordict_unflatten, reconstruct batch_size from the leading dimensions of the actual tensor shapes. The existing jacrev/jacfwd basis-vector mismatch logic is preserved via a backward-compatible "batch_size" legacy path.

Tests

Two new tests added to TestExport in test/test_compile.py:

test_export_dynamic_batch_size_scalar: covers batch_size=x.shape[0] (scalar SymInt, strict=False) — the original issue repro
test_export_dynamic_batch_size_multi_dim: covers batch_size=[b, t] (multi-dim SymInt, strict=True) — the pytree serialization fix

Both tests fail on unpatched code and pass after the fix. The full TestExport suite (5 tests) passes in a clean venv.

Checklist

I have read the CONTRIBUTING.md
Tests added for both bug fixes
Black + flake8 clean
Links to issue [BUG] TensorDict with dynamic, input-dependent batch_size is not torch.export.exportable #1003

…may contain SymInt) During torch.export.export() with dynamic_shapes, TensorDict's _tensordict_flatten stored batch_size (a torch.Size that may contain SymInts) in the pytree context. When torch.export tries to serialize the output pytree spec it calls as_python_constant() on every context value — SymInts are not Python constants, raising AsPythonConstantNotImplementedError. Fix: store batch_dims = len(batch_size) (a plain int, always serializable) instead of batch_size itself. In _tensordict_unflatten, reconstruct batch_size from the first tensor's leading shape dimensions. The existing jacrev/jacfwd basis-vector logic is preserved in the legacy batch_size path for backward compatibility with any serialized tree specs that still carry batch_size. Fixes pytorch#1003

github-actions · 2026-05-16T19:37:32Z

PR Title Label Error

PR title must start with a label prefix in brackets (e.g., [BugFix]).

Current title: fix: store batch_dims (int) in pytree context instead of batch_size (may contain SymInt)

Supported Prefixes

Your PR title must start with exactly one of these prefixes (case-insensitive):

Prefix	Label Applied	Example
`[BugFix]` or `[Fix]`	bug	`[BugFix] Fix memory leak in TensorDict`
`[Feature]`	Feature	`[Feature] Add new storage backend`
`[Doc]` or `[Docs]`	documentation	`[Doc] Update installation guide`
`[Refactor]`	Refactor	`[Refactor] Clean up module imports`
`[CI]`	CI	`[CI] Fix workflow permissions`
`[Test]` or `[Tests]`	Test	`[Test] Add unit tests for nn module`
`[Compile]`	Compile	`[Compile] Fix torch.compile issue`
`[Performance]` or `[Perf]`	Performance	`[Perf] Optimize tensor operations`
`[Deprecation]`	Deprecation	`[Deprecation] Mark old function`
`[Setup]`	setup	`[Setup] Update build configuration`
`[Distributed]` or `[Dist]`	Distributed	`[Distributed] Add scatter collective`
`[Benchmark]` or `[Bench]`	Benchmarks	`[Benchmark] Add compile benchmark`
`[Typing]` or `[Type]`	Typing	`[Typing] Add type stubs`
`[BC-breaking]` or `[BC]`	BC-breaking	`[BC-breaking] Remove deprecated API`
`[Formatting]` or `[Format]`	Formatting	`[Format] Fix code style`
`[Quality]`	Quality	`[Quality] Improve error messages`

Note: Matching is case-insensitive. Common variations (singular/plural) are supported.

_parse_batch_size checked isinstance(batch_size, Number) to handle scalar batch sizes (e.g. batch_size=x.shape[0]), but torch.SymInt does not subclass numbers.Number. During torch.export with dynamic shapes this caused: ValueError: batch size was not specified when creating the TensorDict instance and it could not be retrieved from source. Fix: extend the Number check to also accept torch.SymInt. Also expand the regression test added in the previous commit to cover this scalar SymInt case (the original repro from issue pytorch#1003). Fixes pytorch#1003

… tests

github-actions · 2026-05-17T10:02:05Z

PR Title Label Error

PR title must start with a label prefix in brackets (e.g., [BugFix]).

Current title: Fix dynamic-shape export (#1003): handle scalar SymInt in _parse_batch_size and store batch_dims in pytree context

Supported Prefixes

Your PR title must start with exactly one of these prefixes (case-insensitive):

Prefix	Label Applied	Example
`[BugFix]` or `[Fix]`	bug	`[BugFix] Fix memory leak in TensorDict`
`[Feature]`	Feature	`[Feature] Add new storage backend`
`[Doc]` or `[Docs]`	documentation	`[Doc] Update installation guide`
`[Refactor]`	Refactor	`[Refactor] Clean up module imports`
`[CI]`	CI	`[CI] Fix workflow permissions`
`[Test]` or `[Tests]`	Test	`[Test] Add unit tests for nn module`
`[Compile]`	Compile	`[Compile] Fix torch.compile issue`
`[Performance]` or `[Perf]`	Performance	`[Perf] Optimize tensor operations`
`[Deprecation]`	Deprecation	`[Deprecation] Mark old function`
`[Setup]`	setup	`[Setup] Update build configuration`
`[Distributed]` or `[Dist]`	Distributed	`[Distributed] Add scatter collective`
`[Benchmark]` or `[Bench]`	Benchmarks	`[Benchmark] Add compile benchmark`
`[Typing]` or `[Type]`	Typing	`[Typing] Add type stubs`
`[BC-breaking]` or `[BC]`	BC-breaking	`[BC-breaking] Remove deprecated API`
`[Formatting]` or `[Format]`	Formatting	`[Format] Fix code style`
`[Quality]`	Quality	`[Quality] Improve error messages`

Note: Matching is case-insensitive. Common variations (singular/plural) are supported.

…ch_size for eager The previous approach stored batch_dims unconditionally, breaking torch.func transforms (jacrev, jacfwd, hessian) which run eagerly and rely on detecting batch_size mismatches in _tensordict_unflatten to handle basis-vector shapes. Fix: use is_compiling() to select the context key: - is_compiling() == True → store batch_dims (int): SymInt-safe for torch.export - is_compiling() == False → store batch_size: preserves jacrev mismatch detection The unflatten path is unchanged for eager (batch_size key); the batch_dims path is a simple leading-shape reconstruction used only during compilation where jacrev does not apply.

felixmaximilian · 2026-05-17T10:17:24Z

Hi, the CLA is signed. Could a maintainer please approve the CI run? Happy to address any review feedback.

vmoens

LGTM thanks!

Two fixes are now merged upstream (pytorch/tensordict#1704) but not yet released. This commit applies them locally as a monkey-patch until tensordict ships a release that includes the fix. tensordict_export_patch.apply() must be called before torch.export.export() with dynamic_shapes. The patch is idempotent and matches the upstream fix: - _tensordict_flatten: store batch_dims (int) during compilation instead of batch_size (may contain SymInts); keep batch_size for eager so jacrev/jacfwd mismatch detection is preserved - _parse_batch_size: accept scalar torch.SymInt (e.g. batch_size=x.shape[0]) which does not subclass numbers.Number and previously fell to ValueError episode.py: replace mit.one() set comprehensions with mit.first() + direct .shape access. Creating a Python set forces SymInts to be hashed, which specializes them to concrete values and triggers ConstraintViolationError during dynamic-shape export. test_export.py: remove xfail from test_torch_export_dynamic_shapes; apply the patch within the test. Also add _soft_pending_unbacked session fixture to demote the spurious "pending unbacked symbol u0" warning from tensordict's _device_recorder side-effect. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>

…arse_batch_size and store batch_dims in pytree context (#1704)

Two fixes are now merged upstream (pytorch/tensordict#1704) but not yet released. This commit applies them locally as a monkey-patch until tensordict ships a release that includes the fix. tensordict_export_patch.apply() must be called before torch.export.export() with dynamic_shapes. The patch is idempotent and matches the upstream fix: - _tensordict_flatten: store batch_dims (int) during compilation instead of batch_size (may contain SymInts); keep batch_size for eager so jacrev/jacfwd mismatch detection is preserved - _parse_batch_size: accept scalar torch.SymInt (e.g. batch_size=x.shape[0]) which does not subclass numbers.Number and previously fell to ValueError episode.py: replace mit.one() set comprehensions with mit.first() + direct .shape access. Creating a Python set forces SymInts to be hashed, which specializes them to concrete values and triggers ConstraintViolationError during dynamic-shape export. test_export.py: remove xfail from test_torch_export_dynamic_shapes; apply the patch within the test. Also add _soft_pending_unbacked session fixture to demote the spurious "pending unbacked symbol u0" warning from tensordict's _device_recorder side-effect. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>

@overload

* feat(export): remove EpisodeExport — single Episode path Port of PR #219 ideas onto refactor/timestep-ordering-v2. - Delete EpisodeExport dataclass and pytree registration; EpisodeBuilder.forward now always returns Episode (TensorClass) - Switch TensorDict.from_dict(d, batch_dims=2) → TensorDict(d, batch_size=[b,t]): from_dict goes through from_any → _is_dataclass → hasattr which Dynamo cannot trace - Switch Index.from_dict → Index.from_tensordict(TensorDict(index, batch_size=[t])) for the same reason - Unify PolicyObjective: remove overload/@overload, isinstance(EpisodeExport) branch, and tree_map_with_path path; single Episode._compute_logits - ControlTransformer.forward: remove is_exporting() guard, always return TensorDict - norm.py Scaler/UniformBinner: remove is_exporting()/ValueError split; clamp at output level (Scaler: clamp output to out_range; UniformBinner: clamp bins to [0, bins-1]) rather than clamping the raw input first - test_export.py: episode_export fixture returns Episode; test_episode uses shared ep_dict helper; test_torch_export_fake zips leaves by position (no asdict/is_dataclass); remove policy-continuous tolerance hack (single code path, bit-exact now); add test_onnx_inference regression test (ORT vs eager) - Add CLAUDE.md with export architecture documentation Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com> * fix: guard and test dynamic-shape export limitation (tensordict#1003) - Raise NotImplementedError with a clear message in EpisodeBuilder.forward when torch.export detects a dynamic batch dim (SymInt), rather than letting a cryptic tensordict traceback propagate - Add xfail test_torch_export_dynamic_shapes: canary that will flip to xpass once pytorch/tensordict#1003 is resolved Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com> * fix: apply tensordict#1003 workaround for dynamic-shape export Two fixes are now merged upstream (pytorch/tensordict#1704) but not yet released. This commit applies them locally as a monkey-patch until tensordict ships a release that includes the fix. tensordict_export_patch.apply() must be called before torch.export.export() with dynamic_shapes. The patch is idempotent and matches the upstream fix: - _tensordict_flatten: store batch_dims (int) during compilation instead of batch_size (may contain SymInts); keep batch_size for eager so jacrev/jacfwd mismatch detection is preserved - _parse_batch_size: accept scalar torch.SymInt (e.g. batch_size=x.shape[0]) which does not subclass numbers.Number and previously fell to ValueError episode.py: replace mit.one() set comprehensions with mit.first() + direct .shape access. Creating a Python set forces SymInts to be hashed, which specializes them to concrete values and triggers ConstraintViolationError during dynamic-shape export. test_export.py: remove xfail from test_torch_export_dynamic_shapes; apply the patch within the test. Also add _soft_pending_unbacked session fixture to demote the spurious "pending unbacked symbol u0" warning from tensordict's _device_recorder side-effect. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com> * fix: restore unresolved-attribute ty:ignore in mask.py (ty@latest regression) * chore: add CLAUDE.md, .DS_Store, .vscode/ to .gitignore These are local developer artifacts that should not be tracked: - CLAUDE.md: AI assistant codebase notes (editor/project-local) - .DS_Store: macOS Finder metadata - .vscode/: editor config (launch.json etc) — will be added via a dedicated PR directly to main Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com> * fix: remove extra blank line in episode.py (ruff-format) Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com> * fix: use TensorDict() constructor for embeddings to avoid dynamo hasattr failure TensorDict.from_dict internally calls from_any → is_dataclass → hasattr, which dynamo cannot trace during strict torch.export. Use the direct TensorDict(embeddings, batch_size=[b, t]) constructor instead, consistent with the other fields in this constructor call. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com> * refactor: full TensorDict migration in EpisodeBuilder.forward After self.projections(), wrap projected_embeddings into a TensorDict and use it throughout the rest of forward(). This eliminates pytree API usage in _build_index, _build_role_embeddings, and the embedding sum / pack loop: - _build_index: .batch_size[1] + .get((mod, name)) instead of mit.first / tree_leaves / key_get / MappingKey - _build_role_embeddings: .batch_size[1] instead of mit.first / tree_leaves, and the None-handling branch becomes a torch.zeros_like fallback since filter_non_tensor_data() already removed None entries - Embedding sum: projected_td.apply(torch.add, role_td) instead of tree_map(lambda p, r: p + r, ...). Note: TensorDict's __add__ uses _foreach_add which the ONNX exporter cannot convert; .apply(torch.add) produces per-tensor torch.add calls that ONNX handles fine. - Pack: embeddings.get(k) instead of key_get + MappingKey Device is passed explicitly to the helpers since TensorDict.device is None unless set in the constructor, and adding device=device to the constructor triggers a "Can only mark one TensorDict at a time." error in dynamo when tracing nested dicts during export. key_get is no longer used → dropped from imports. mit.first / tree_leaves are still used once each in forward() for the initial b/t/device extraction from input_tokens (a pytree before wrapping). Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com> * refactor: replace mit.first/tree_map with next()/dict comprehension Drop the `more_itertools` import — `mit.first(...)` is just `next(...)` for the b/t/device extraction. Replace `tree_map(special_tokens, is_leaf=...)` with an explicit 2-level dict comprehension; clearer and avoids the obscure `is_leaf=isinstance tuple` pytree contract. Still uses `tree_leaves` once to probe input_tokens for b/t/device. To eliminate that entirely we'd need `batch` to be a TensorDict (with batch_size=[b, t]) upstream — that's a dataloader/Lightning entry-point change tracked separately. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com> * refactor(tests): simplify test_episode to direct assert_close Episode is now a TensorClass on both eager and export paths (single Episode return after this PR's collapse). assert_close handles TensorClass comparison directly — no need to flatten via pytree and key_get into per-leaf assertions, no need for a custom ep_dict adapter that re-wraps embeddings/embeddings_flattened. 13 lines → 1. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com> --------- Co-authored-by: Max Moeller <max@yaak.ai> Co-authored-by: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>

meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label May 16, 2026

github-actions Bot added Test Compile torch.compile related labels May 16, 2026

felixmaximilian changed the title ~~fix: store batch_dims (int) in pytree context instead of batch_size (may contain SymInt)~~ [BugFix]: store batch_dims (int) in pytree context instead of batch_size (may contain SymInt) May 16, 2026

github-actions Bot added the bug Something isn't working label May 16, 2026

felixmaximilian marked this pull request as draft May 17, 2026 09:55

Max Moeller added 2 commits May 17, 2026 11:56

refactor(test): split dynamic batch_size export test into two focused…

0733453

… tests

felixmaximilian changed the title ~~[BugFix]: store batch_dims (int) in pytree context instead of batch_size (may contain SymInt)~~ Fix dynamic-shape export (#1003): handle scalar SymInt in _parse_batch_size and store batch_dims in pytree context May 17, 2026

felixmaximilian changed the title ~~Fix dynamic-shape export (#1003): handle scalar SymInt in _parse_batch_size and store batch_dims in pytree context~~ [BugFix] Fix dynamic-shape export (#1003): handle scalar SymInt in _parse_batch_size and store batch_dims in pytree context May 17, 2026

felixmaximilian marked this pull request as ready for review May 17, 2026 10:17

vmoens approved these changes May 17, 2026

View reviewed changes

vmoens merged commit babcd6e into pytorch:main May 17, 2026
66 of 69 checks passed

egorchakov mentioned this pull request May 18, 2026

feat(export): remove EpisodeExport — single Episode path yaak-ai/rmind#231

Merged

1 task

vmoens pushed a commit that referenced this pull request May 21, 2026

[BugFix] Fix dynamic-shape export (#1003): handle scalar SymInt in _p…

664bad7

…arse_batch_size and store batch_dims in pytree context (#1704)

vmoens pushed a commit that referenced this pull request May 21, 2026

[BugFix] Fix dynamic-shape export (#1003): handle scalar SymInt in _p…

9fb10b3

…arse_batch_size and store batch_dims in pytree context (#1704)

Navigation Menu

Uh oh!

[BugFix] Fix dynamic-shape export (#1003): handle scalar SymInt in _parse_batch_size and store batch_dims in pytree context - #1704

[BugFix] Fix dynamic-shape export (#1003): handle scalar SymInt in _parse_batch_size and store batch_dims in pytree context#1704
vmoens merged 4 commits into
pytorch:mainfrom
felixmaximilian:fix/dynamic-shape-export-1003

felixmaximilian commented May 16, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented May 16, 2026

Uh oh!

github-actions Bot commented May 17, 2026

Uh oh!

felixmaximilian commented May 17, 2026

Uh oh!

vmoens left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

felixmaximilian commented May 16, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Problem

Fix

Tests

Checklist

Uh oh!

github-actions Bot commented May 16, 2026

PR Title Label Error

Supported Prefixes

Uh oh!

github-actions Bot commented May 17, 2026

PR Title Label Error

Supported Prefixes

Uh oh!

felixmaximilian commented May 17, 2026

Uh oh!

vmoens left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

felixmaximilian commented May 16, 2026 •

edited

Loading