[zingo]

Serialize bundled inputs and expected outputs in the dimension order declared by the program. This preserves channels-last values and avoids layout mismatches when eager test tensors have permuted strides.

Reconstruct ATen tensors using the serialized dimension order. Keep layout conversion within BundledProgram instead of adding a preserve_layout flag to TensorSpec.from_tensor.

Cover channels-last tensors, offset and non-dense views, and permuted outputs, including portable runtime verification. Correct the dim_order_from_stride return annotation and update Buck dependencies.

Assisted by OpenAI Codex.

cc @digantdesai @freddan80 @per @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/23927

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• PRs based before #202327 fail all Linux builds: can't find build-external-packages action

⏳ No Failures, 335 Pending

As of commit 3321612 with merge base 403800d ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[zingo]

@claude - Review this PR.

[claude]

Claude finished @zingo's task in 2m 17s —— View job

---

The serialization logic looks correct; I found one performance regression on the pte_file_path path, two smaller robustness and cleanup issues, and no blockers.

Detailed review

PR Review: #23927

Summary

This PR serializes bundled inputs and expected outputs in the dim order each method declares for its I/O, and rebuilds ATen bundled tensors with matching strides. The conversion is correct, but serialize_to_schema() now fully deserializes the PTE when pte_file_path is used, which costs a lot on large models.

Performance

• devtools/bundled_program/core.py:99-102: serialize_to_schema() now always calls _deserialize_pte_binary(program_bytes) when the bundle is built from pte_file_path. That function does two expensive things:

  • It runs the whole flatbuffer through a flatc JSON round trip (_program_flatbuffer_to_json), which includes any inline constant_buffer data.
  • It copies every segment into new bytes objects in _restore_segments, including delegate blobs and constant/named data.

  The only things used afterwards are plan.inputs, plan.outputs and the tensor dim_order of those values. Before this PR, the file path only read the raw bytes. For multi-GB PTEs (LLMs, large delegate payloads) this is a large time and peak-memory regression just to read a few dim_order bytes. Two options: parse only the flatbuffer part, slicing to eh.program_size from the extended header so _restore_segments is skipped (this still pays the JSON cost for inline constants), or read the execution plan with the generated flatbuffer Python accessors. Also consider caching the parsed program so the work isn't repeated. Fix this →

Code Quality

• devtools/bundled_program/core.py:122-144: With pte_file_path, _assert_valid_bundle is still skipped (core.py:73-74), but the program is now parsed anyway. A mismatched test case now fails with raw errors instead of the existing clear assertion messages:

  • too many inputs or outputs gives an IndexError on plan.inputs[input_idx] or plan.outputs[output_idx];
  • a tensor passed where the program expects Int gives an AttributeError on .val.dim_order;
  • a tensor of the wrong rank gives an error from permute.

  Since program is available on both paths now, run the same _assert_valid_bundle checks (refactored to take a Program) for the file path too.

• exir/tensor.py:208-215: The only change to TensorSpec.from_tensor is reflowing its signature onto several lines and deleting the "for non-contiguous tensors" comment. Behaviour is unchanged and the edit looks left over from the dropped preserve_layout approach. Please revert it to keep the diff focused.

Recommendation

Approve

Nothing blocks merging. I'd still address the full-PTE deserialization before landing, because it penalizes every file-based bundle in proportion to model size.

[zingo]

@claude - re-Review this PR.

[claude]

Claude finished @zingo's task in 3m 38s —— View job

---

One issue from my first review is still open: bundles built from pte_file_path still skip validation, so a bad test suite crashes with a bare Python error. I also found two small problems and nothing that blocks merging.

Detailed review

PR Review: #23927

Summary

This PR serializes bundled inputs and expected outputs in the dim order each method declares for its I/O, and rebuilds ATen bundled tensors with the matching strides. Two findings from the earlier review are resolved: the full-PTE deserialization (the file path now reads only I/O metadata from the flatbuffer) and the leftover TensorSpec.from_tensor change. What remains is missing validation on the pte_file_path path, a type annotation fix that misses two declarations, and a dependency on private serializer helpers.

Code Quality

• devtools/bundled_program/core.py:119, :131, :147: _assert_valid_bundle still runs only when executorch_program is given (core.py:83-84). With pte_file_path, serialize_to_schema() now looks up program metadata without any checks. A bad test suite fails with a bare Python error instead of a clear assertion message:

  • a method name that is not in the PTE gives KeyError on method_io[method_test_suite.method_name];
  • more test inputs or outputs than the plan has gives IndexError on plan.inputs[input_idx] or plan.outputs[output_idx];
  • a tensor passed where the program expects Int, Bool or Double gives AttributeError: ... has no attribute 'dim_order';
  • a tensor of the wrong rank gives a RuntimeError from tensor.permute(dim_order) at core.py:179.

  _get_method_io now gives both paths the same per-method I/O list, so the method-name, input/output-count, type and rank checks can run on that list for both paths. There are no tests for these error cases on the file path. Fix this →

• exir/tensor.py:51-53: dim_order_from_stride now returns Tuple[int, ...] and the typing.cast is gone. Two attributes that store its result are still declared Tuple[bytes]:

  • exir/tensor.py:185: self.dim_order: Tuple[bytes] = dim_order_from_stride(self.stride);
  • exir/passes/replace_view_copy_with_view_pass.py:318: the same assignment.

  The cast hid this mismatch, so internal Pyre should now report it. Update both declarations in this PR. Fix this →

• devtools/bundled_program/core.py:22-27, :216-237 and devtools/bundled_program/BUCK:28: Low-severity design concern. The file path now relies on the private helpers _patch_generated_flatbuffer_aliases and _convert_from_flatbuffer_dataclass and on generated EValueT/Program, and it adds a dependency on //executorch/exir/_serialize:lib.

  • The comment on that target asks people to check before adding users, because of ODR concerns with other Flatbuffers libraries.
  • Every I/O EValue is converted in full to a core_schema dataclass, yet only dim_order is used. Reading the generated Tensor's dim order directly, or adding a small public helper in exir/_serialize, would reduce coupling to serializer internals.
  • The exir/_serialize owners should sign off on the new dependency.

Recommendation

Approve

Nothing blocks merging. The pte_file_path validation gap is worth fixing before landing, because file-based users now get unclear crashes where they used to get a bundle that was silently wrong.

[zingo]

@claude please re-review I have a follow up PR for "BundleIO: Validate file-based and in-memory test data consistently

Share validation using method I/O metadata for both program sources.
Check method names, value counts, types, tensor dtypes and ranks before
serialization while keeping expected outputs optional.

The runtime verifier expects one reference value per model output.
Require all expected outputs, or none for profiling. Reject partial
lists when bundling to avoid out-of-range access during verification.

Add regression coverage for invalid test data and preserve support for
tensor sizes that vary within the same rank."
so it can be reviewed handled differently but it could also be incoperated to this PR. Let me know.

[claude]

Claude finished @zingo's task in 5m 0s —— View job

---

I'd keep the validation work as a separate follow-up PR. This PR has no blockers; three of its new tests pass on the old code too, and one earlier design concern is still open.

On the follow-up PR

Recommendation: keep it separate and land it right after this one.

• This PR stands alone. It fixes the layout and dim-order handling.
• The follow-up changes behaviour for users. It rejects test cases that give only some of the expected outputs. Callers who do that today will hit an assertion they didn't hit before. A change like that deserves its own review and its own release notes: devtools entry, rather than being buried in a layout fix.
• Merging them would slow this PR down. All three reviewers found nothing blocking here, and adding the validation would start another review cycle.
• Land the follow-up soon after. This PR adds new metadata lookups on the pte_file_path path, so until the follow-up merges, a bad test suite crashes with a bare KeyError or IndexError. Stacking the follow-up on this branch, or opening it straight after, keeps that window short.

Detailed review

PR Review: #23927

Summary

This PR serializes bundled inputs and expected outputs in the dim order each method declares, and rebuilds ATen bundled tensors with matching strides. Most points from the previous review are fixed: full-PTE deserialization, the leftover TensorSpec change, and the stale Tuple[bytes] annotations. What remains is minor: three tests that don't catch the regression they're named for, an unclear crash when dim_order is missing, and the earlier concern about depending on serializer internals. Missing validation on the pte_file_path path is left out here because the follow-up PR covers it.

Testing

• Three new tests pass on the old code as well. The old path called TensorSpec.from_tensor(t, const=True), which runs tensor.contiguous() (exir/tensor.py:209-212). So for any program whose I/O is contiguous, the old code already wrote dim_order=[0,1,…] and contiguous bytes.

  • devtools/bundled_program/test/test_bundle_data.py:148 test_bundle_uses_program_dim_order: the Permute program's input and output are contiguous, so the "program dim order" equals what the old code produced. A transposed input or permuted expected output doesn't change that.
  • test_bundle_data.py:106 test_non_dense_view_uses_only_its_own_elements: .contiguous() on base[:, :, ::2, :] already gives a new dense tensor of exactly view.nbytes.
  • extension/pybindings/test/test_pybindings.py:126 test_bundled_permute_output: permute_copy returns a contiguous output, which matches the old contiguous serialization of expected.

  These are fine as extra coverage, but the PR description lists "non-dense views" and "permuted outputs" among the fixes. Either reword the description, or make the tests fail without the change. For example, use a program whose I/O dim order is channels-last and feed it tensors that are not channels-last. The channels-last tests, the storage-offset test, the pending-negative test and test_bundled_channels_last_input_and_output do fail on the old code.

Code Quality

• devtools/bundled_program/core.py:179: if a program's I/O tensor has no dim_order vector, the call becomes tensor.permute(None) and raises a bare TypeError.
  • On the file path the generated TensorT leaves dimOrder as None when the field is absent, and _convert_from_flatbuffer_value passes that through.
  • The current emitter always writes dim_order, so only older or third-party PTEs would hit this. Low confidence that this matters in practice.
  • A clear assertion, or a fallback to list(range(tensor.dim())), would fix it. This isn't rank or type validation, so the follow-up may not cover it.
• devtools/bundled_program/core.py:22-27 and devtools/bundled_program/BUCK:28: still open from the previous review.
  • core.py imports the private _patch_generated_flatbuffer_aliases and _convert_from_flatbuffer_dataclass, plus the generated Program/EValueT classes.
  • It adds a dependency on //executorch/exir/_serialize:lib, whose BUCK comment (exir/_serialize/BUCK:61-63) says "Please ask before changing this" because of the ODR risk.
  • _patch_generated_flatbuffer_aliases() also re-imports the whole generated package on every file-based serialize_to_schema() call.
  • Please get sign-off from the exir/_serialize owners, or expose a small public helper there that returns the per-method I/O dim orders.

Recommendation

Approve

Nothing blocks merging. The test findings are about how well the tests catch regressions, not about the fix being wrong. Land the validation follow-up soon after this PR so the period where pte_file_path users get unclear crashes stays short.

Not checked at runtime: whether test_bundled_channels_last_input_and_output gives the expected output strides under the aten_lib pybindings target. The reviews only read the code and ran no tests, because commands outside a small allowlist needed approval in this environment.

[zingo]

Hi @rascani and @digantdesai would this need internal testing?