[shoumikhin]

Problem

Some LeRobot policies fail with the MLX backend (Apple's MLX delegate):

• smolvla, smolvla_libero and pi0 stop inside the MLX partitioner with
  SpecViolationError: Operator '<EdgeOpOverload: aten.clone.default> ...' is not an allowed operator type.
• vla_jepa, with Fix export of zero-stride constants and inputs, and of constants that delegates share #24404 applied, lowers, but then to_executorch fails with
  AttributeError: `_prop_tensor_constant551` is not a buffer.
• vla_jepa and lawam lower to MLX, then fail at run time with
  [Gather] Boolean indices not supported.

The clone error. When constant_prop_pass folds a program output into a constant, it puts a clone in front of that output. The clone was always an Edge dialect op, even when the pass ran on an ATen program before to_edge. The MLX partitioner checks the ATen program in transform_for_pre_decomposition, and the Edge op fails that check.

The "not a buffer" error. When ops in a delegate and ops outside it read the same constant, duplicate_constant_node gives every user but one its own copy of the placeholder. A program that lists a lifted constant twice cannot be unlifted, so to_executorch fails when it retraces the delegate call. Core ML hits the same bug earlier, as KeyError: '<buffer>' in coremltools.

Fix

• constant_prop_pass checks the program's dialect. An ATen or training program now gets torch.ops.aten.clone.default, which to_edge later turns into the right Edge op.
• duplicate_constant_node groups the users of a constant by delegation tag. Users with the same tag, or with none, share one placeholder, and only the other groups get a copy. Tests that call it directly now tag their users, since untagged users no longer get copies.
• The MLX runtime gathers only with integer indices. The MLX index handler now replaces a boolean mask that is a constant of the program (for example one folded by constant_prop_pass) with its integer positions. Any other boolean index stays on the CPU.

Testing

• Each fix has a new test that fails without it.
• The exir pass, emit and to_backend test files have no new failures.
• Small MLX repros of each error fail before and pass after. A Core ML repro of the shared buffer case raised the KeyError before and lowers after.
• vla_jepa also reads a shared constant through a list (in index.Tensor), a separate bug that Fix export of zero-stride constants and inputs, and of constants that delegates share #24404 fixes. With both, the full vla_jepa graph gets through MLX lowering and to_executorch.

#24404 also edits duplicate_constant_node, so whichever merges second needs a small rebase.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

⏳ 5 Pending, 1 Unrelated Failure

As of commit 052ecdd with merge base a489dca ():

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / nxp-mcuxpresso-test / linux-job (gh) (trunk failure)
  undefined reference to SYS_INFO_PrintAll'`

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

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[shoumikhin]

@claude /pr-review

[claude]

Claude finished @shoumikhin's task in 19m 19s —— View job

---

Summary: Nothing blocks merging. The findings are three test gaps (the new per-tag grouping in duplicate_constant_node, where the new test lives, and the ATen clone after to_edge), a note for code that calls duplicate_constant_node directly, and one optional simplification.

Detailed review

PR Review: #24481

Summary

The PR does two things:

• constant_prop_pass now puts an aten.clone.default in front of folded outputs when the program is not Edge yet.
• duplicate_constant_node now makes one placeholder per delegation-tag group instead of one per user.

Nothing blocks merging. The issues below are test gaps, one note for direct callers, and one optional simplification.

Code Quality

• exir/backend/canonical_partitioners/duplicate_constant_node_pass.py:112-132: this PR rewrites the whole args/kwargs rebuild anyway, so it could become one call: user.replace_input_with(constant_or_attribute_node, copied_constant_or_attribute_node). That call is already used in exir/passes/quant_fusion_pass.py. It also rewires inputs nested in lists, such as index.Tensor indices, which the direct arg != constant_or_attribute_node comparison misses. That nested case is the one the description leaves to Fix export of zero-stride constants and inputs, and of constants that delegates share #24404, so agree with Fix export of zero-stride constants and inputs, and of constants that delegates share #24404 on which PR makes this change. Fix this →

Testing

• The grouping is tested in only one order.

  • exir/backend/test/test_utils.py:402-453 always has a delegated op as the constant's first user: add (tag0) in the first to_backend and sub (tag1) in the second. So the delegate always keeps the original placeholder.
  • The other order is never run. There, a CPU op comes first, the original stays at top level, _assign_new_tag (exir/backend/utils.py:205) re-tags it None, and the delegate gets _copy_0. A second case with e.g. torch.sin(self.const) first in forward would cover it. The node names in TagNodes would shift.
  • No test checks the new duplicate_constant_node rule itself. The edited tests (exir/backend/test/test_passes.py:41-43, 84-86, backends/transforms/test/test_create_delete_constant_placeholder.py:214-217) give every user its own tag, so they only re-test one copy per user. Worth adding: same-tag users and untagged users get set(), and three groups get _copy_0 and _copy_1.

  Fix this →

• The new to_backend test sits apart from the tests like it.

  • exir/backend/test/test_utils.py:402: the existing constant-splitting tests, test_partitioner_splits_constant_data and test_partitioner_alert_split_constant_data, are in exir/backend/test/test_partitioner.py. They use the same ExecutorBackend.
  • That BUCK target already depends on executor_backend_preprocess and preloads executor_backend_register.
  • Moved there, the test would not need the new dep at exir/backend/test/BUCK:381. It could also run the lowered program with _load_for_executorch_from_buffer, as test_partitioner_splits_constant_data does, and compare against eager. Right now it only checks that to_executorch() does not raise.

  Fix this →

• The ATen clone is not followed through to_edge.

  • exir/tests/test_passes.py:1640-1660 checks the clone target and _validate() on the ATen and training programs.
  • It does not check what the description says happens next: that to_edge turns aten.clone.default into the right Edge op. That is the path the MLX export takes after transform_for_pre_decomposition.
  • A case in exir/tests/test_partitioner_pre_decomposition.py would cover the reported flow end to end: transform_for_pre_decomposition calls constant_prop_pass, then to_edge_transform_and_lower runs, and the test asserts _clone_dim_order (with the default EdgeCompileConfig) feeds the folded output.

  Fix this →

Backward Compatibility

• exir/backend/canonical_partitioners/duplicate_constant_node_pass.py:93: duplicate_constant_node has no underscore, so outside code can call it.
  • When its users are untagged or all share one tag, it now makes no copies and returns set(). Before, every user after the first got a copy.
  • The new docstring line describes this. The only in-tree production caller, _maybe_duplicate_constant_nodes (exir/backend/utils.py:242), runs after tagging, so it is unaffected.
  • Out-of-tree code that calls it before tagging, as the three in-tree tests did before this PR changed them, will now get no copies and no warning.
  • The PR has the release notes: exir label, so mention this in the release note.

Recommendation

Approve. Nothing blocks. The tests above are worth adding in this PR, most of all the case where a CPU op is the first user, because that is the path that re-tags the original placeholder.

I did not run any tests: this job can't run Python. Everything above comes from reading the code.

---

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[shoumikhin]

Thanks for the review. The latest push changes 2 tests:

• The shared-constant test now also runs with the CPU ops reading the constant first. Then the original stays in the top-level program and the delegate gets the copy.
• The dialect test now checks that to_edge turns the ATen clone into _clone_dim_order.

What I left as is:

• replace_input_with. It would also rewire a constant read through a list. Fix export of zero-stride constants and inputs, and of constants that delegates share #24404 already makes that change in this function, so I kept it out here, so the 2 PRs do not make the same change.
• Moving the test to test_partitioner.py. pytest.ini skips that file, so GitHub CI would stop running the test. test_utils.py is run there. The bug is raised in to_executorch, so the test checks up to that point.
• A separate unit test for the grouping rule. I tried breaking the grouping in a few ways: no grouping, one group for all users, a copy for every group, and only the first user of a group rewired. Each one makes the new test or the edited direct tests fail.

About direct callers: a graph with no delegation tags now gets no copies. The docstring, the commit message and the PR description say this, so it will show up in the release notes. The only non-test caller in the tree runs after tagging.

[shoumikhin]

@claude /pr-review

[claude]

Claude finished @shoumikhin's task in 17m 26s —— View job

---

Summary: Nothing blocks merging. There are three minor findings: one subtest of the new to_backend test tags the wrong add, GitHub CI skips the file that holds the new constant_prop_pass test, and the pass's default skip list still matches only Edge ops on the ATen programs this PR now handles (that last one predates the PR).

Detailed review

PR Review: #24481

Summary

The PR makes two changes:

• constant_prop_pass now puts an ATen clone in front of folded outputs when the program is not Edge yet.
• duplicate_constant_node now gives each delegation tag one placeholder, instead of one per user.

Nothing blocks merging. The findings are two test gaps and an Edge-only default in the same pass. That default predates the PR.

Code Quality

• The default skip list only matches Edge ops, so it skips nothing on the ATen programs this PR now handles.

  • exir/passes/constant_prop_pass.py:38-43 builds _DEFAULT_SKIP_TARGETS from exir_ops.edge.aten.full.default and the Edge forms of the quant primitives.
  • EdgeOpOverload (exir/dialects/edge/_ops.py:291) defines no __eq__ and no __hash__. So the check at :262 never matches torch.ops.aten.full.default or the quantized_decomposed ops, and every op with constant inputs gets run and folded (:281-304).
  • What that means for a run before to_edge:
    • It folds aten.full, which is the size growth the comment at :36 warns about.
    • On a PT2E-quantized model it also folds dequantize_per_* of a constant int8 weight, so the weight is stored as float.
  • This predates the PR, so a follow-up is fine.
  • Fix: choose the default by dialect too, i.e. torch.ops.aten.full.default plus _QUANT_PRIMITIVES when the program is not Edge. The custom_skip_targets: Optional[set[EdgeOpOverload]] type hint makes the same assumption.

  Fix this →

Testing

• With cpu_first=True, the first TagNodes tags the wrong add.

  • exir/backend/test/test_utils.py:411-412 traces torch.sin(self.const) + torch.cos(self.const) first. So that add is named aten_add_tensor, and x + self.const becomes aten_add_tensor_1.
  • :453 then puts the sin+cos add, which does not read the constant, into tag0 along with mul. That delegate reads the constant once, not twice as the comment at :445-446 says, and x + self.const stays on the CPU.
  • The gap is small: the second to_backend (sub and div) still covers a delegate whose two readers share the copy.
  • Fix: return torch.sin(self.const) and torch.cos(self.const) as separate outputs. That removes the second add, and the same tags then work in both subtests.

  Fix this →

• GitHub CI does not run the new constant_prop_pass test.

  • The test is at exir/tests/test_passes.py:1640. pytest.ini:50 has --ignore=exir/tests/test_passes.py, and no script in .ci/ or .github/ names that file, so only Buck runs it.
  • The reply gave this same reason for keeping the to_backend test out of test_partitioner.py.
  • If the clone fix should also run in GitHub CI, the test needs a file pytest collects. For example, a new exir/tests/test_constant_prop_pass.py with a target in exir/tests/targets.bzl.

  Fix this →

Recommendation

Approve. Nothing blocks. The cpu_first tagging is a quick fix, and the skip-list default can wait for a follow-up.

I did not run any tests, because this job can't run Python. Everything above comes from reading the code.

[shoumikhin]

Thanks, good catch on the test. The latest push fixes it:

• Wrong add in the CPU-first case. sin and cos are now separate outputs, so there is no second add. In both cases the first delegate now holds add and mul, and the second holds sub and div. Each order also catches a mistake the other misses: copying for tagged groups only, or for the untagged group only.

What I left as is:

• The constant_prop_pass test runs only in Buck. All the other constant_prop_pass tests are in the same file, so I kept it with them. The difference from the to_backend test is that moving that one would have taken away GitHub coverage it already has. Moving this one into a new file would add coverage for one test only, and that seems better done for the whole group at once.
• The default skip list only matches Edge ops. I agree it is wrong for ATen and training programs, and it is older than this PR. It changes which ops get folded, which is a separate behavior change, so I would rather fix it in its own PR.

[shoumikhin]

@claude /pr-review

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[claude]

Claude finished @shoumikhin's task in 16m 19s —— View job

---

Summary: Nothing blocks merging. There are two minor findings: no test has one to_backend call split a shared constant three ways, and the new clone choice relies on the verifier's dialect label, which Vulkan sets to TRAINING on Edge graphs (latent).

Detailed review

PR Review: #24481

Summary

The PR makes constant_prop_pass give folded outputs of non-Edge programs an ATen clone, and makes duplicate_constant_node give each delegation-tag group one placeholder. Nothing blocks merging. There is one test gap and one latent backward-compatibility risk.

Testing

• No test has one to_backend call split a shared constant three ways.

  • exir/backend/test/test_utils.py:452-459: each TagNodes has one tag. So every duplicate_constant_node call in the test sees exactly two groups: one delegate and the untagged ops. The direct tests this PR edits each give the constant two users, so they also stop at two groups.
  • What never runs is one partitioner returning two partitions that both read the constant, next to CPU users. That happens when an unsupported op splits a graph into several partitions that read the same constant. It is also the only case where one call makes both _copy_0 and _copy_1.
  • So changing list(groups.values())[1:] to [1:2] at exir/backend/canonical_partitioners/duplicate_constant_node_pass.py:46, which copies for the second group only, passes the new test and the edited direct tests. That change is not one of the mutations listed in the reply.
  • Fix: add one more case to the same test, using a single TagNodes({"aten_add_tensor": "tag0", "aten_mul_tensor": "tag0", "aten_sub_tensor": "tag1", "aten_div_tensor": "tag1"}). It should assert [1, 1, 1] and then call to_executorch().

  Fix this →

Backward Compatibility

• The clone choice relies on the verifier's dialect label, and some Edge graphs carry a different label.

  • exir/passes/constant_prop_pass.py:466 gives an ATen clone to every program whose dialect is not "EDGE". dialect comes from the program's verifier, not from the ops in its graph.
  • Vulkan's preprocess runs each ExportPass with override_verifiers=[_any_op] (backends/vulkan/vulkan_preprocess.py:85). _any_op sets dialect = "TRAINING" only to skip the functional check (:63-65). A plain callable later in the same pass list receives the relabeled program (:91-92).
  • No in-tree code calls constant_prop_pass on such a program, so nothing in the tree changes. But a backend that adds it to a pass list run this way (Vulkan, or an out-of-tree backend using the same pattern) now gets torch.ops.aten.clone.default inside an Edge graph, where it used to get the Edge clone.
  • Keeping the label as the signal is the simple choice. If you keep it, add one line to the constant_prop_pass docstring (:495-505) saying the clone follows exported_program.dialect, so callers can see the dependency.

  Fix this →

Recommendation

Approve. Nothing blocks. The test case is a few lines in the existing test. The dialect risk is latent, so a docstring line is enough.

I did not run any tests, because this job can't run Python. Everything above comes from reading the code.

[shoumikhin]

Thanks. Both points are fixed in the latest push:

• One partitioner with 2 delegates. The shared-constant test has a new case where one to_backend call takes both delegates. In that case a single call to duplicate_constant_node makes 2 copies. Copying for the second group only now fails that case.
• The clone follows the dialect label. The constant_prop_pass docstring now says which clone a constant output gets, and that the choice follows exported_program.dialect.

[shoumikhin]

@claude /pr-review

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[claude]

Claude finished @shoumikhin's task in 14m 9s —— View job

---

Summary: This round found no new issues, and nothing blocks merging.

Detailed review

PR Review: #24481

Summary

The PR does two things:

• constant_prop_pass now puts an ATen clone in front of a folded output when the program is not Edge yet.
• duplicate_constant_node now gives each delegation-tag group one placeholder, instead of one per user.

This round found no new issues. Two points from earlier rounds came up again and are left out: choosing the clone by exported_program.dialect, which the docstring now documents, and a separate unit test for the grouping rule. Both were already answered in this thread.

Recommendation

Approve. Nothing blocks.

I did not run any tests, because this job can't run Python. Everything above comes from reading the code.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota. Try requesting a review at a lower effort level.

[shoumikhin]

@claude /pr-review

[claude]

Claude finished @shoumikhin's task in 24m 41s —— View job

---

Summary: Nothing blocks merging. There are two minor findings: an MLX builder comment still says shared constants are copied once per user, and the new constant-mask lowering makes one int64 index per mask dimension, so a multi-dimensional mask goes through a gather path that refuses dynamic dimensions.

Detailed review

PR Review: #24481

Summary

The PR makes three changes:

• constant_prop_pass now puts an ATen clone in front of the folded outputs of a non-Edge program.
• duplicate_constant_node now gives each delegation-tag group one placeholder.
• The MLX index handler now replaces a constant boolean mask with its integer positions.

Nothing blocks merging. Both findings are minor, and both involve the MLX side. The mask handler (052ecdd) is new since the last round.

Code Quality

• A comment in the MLX builder still says shared constants are copied once per user.

  • backends/mlx/builder/program_builder.py:529-530 (not in the diff) says "A constant shared with other partitions is duplicated once per user, so several placeholders here can name the same target." That is the stated reason for the target_to_constant_slot dedupe at :523-537.
  • After this PR, duplicate_constant_node makes one copy per delegation tag (exir/backend/canonical_partitioners/duplicate_constant_node_pass.py:36-46). _assign_new_tag (exir/backend/utils.py:198-206) then gives each copy the tag of its users. So the users of a constant in one partition now share one placeholder, and the comment's reason no longer holds.
  • Fix: either say in the comment what can still give one partition several placeholders with the same target, or remove the dedupe if nothing can. Fix export of zero-stride constants and inputs, and of constants that delegates share #24404 changes how reads through a list are rewired, so it may change that answer.

  Fix this →

Performance

• A constant mask becomes one int64 index per mask dimension.

  • backends/mlx/ops.py:2448-2451: value.nonzero(as_tuple=True) returns one int64 tensor per mask dimension, and make_or_get_constant (backends/mlx/builder/program_builder.py:260-278) stores each one unchanged. So a d-dimensional mask with K true entries is stored as d·K int64 values. For example, a mostly-true 1024×1024 mask becomes about 16 MB of positions, where the mask itself is 1 MB.
  • With d ≥ 2, the indices go through the multi-index GatherNode path (ops.py:2508-2530). That path raises for any non-indexed dimension with a dynamic size. So a 2-D mask is not lowered when another input dimension is dynamic, for example x[:, mask2d] with a dynamic batch dimension.
  • Fix, if large masks or dynamic shapes are expected: the masked dimensions are static, because they must match the mask. Reshape them into one dimension, then gather once through the single-index TakeNode path with value.flatten().nonzero().squeeze(1) as the index. That stores K values, and that path does not require static sizes.

  Fix this →

Recommendation

Approve. Nothing blocks. The comment is a quick fix. The mask layout can wait unless the target models have large multi-dimensional masks or dynamic shapes.

I did not run any tests, because this job can't run Python. Everything above comes from reading the code.