Support non-tensor attrs in __tensor_flatten__ by aorenste · Pull Request #176457 · pytorch/pytorch

aorenste · 2026-03-04T18:13:54Z

Stack from ghstack (oldest at bottom):

-> Support non-tensor attrs in __tensor_flatten__ #176457

The first return value of __tensor_flatten__ can now contain opaque
(non-tensor) values such as DeviceMesh. Opaques flow through the flat
arg list alongside tensors: they get indices, become graph inputs/outputs,
and are pulled from all_args by position during subclass reconstruction
— the same as PlainTensorMeta. An empty OpaqueMeta marker in
SubclassCreationMeta.attrs distinguishes opaque slots from tensor slots
(needed by process_runtime_tangent, which skips non-differentiable
opaques).

All code that previously iterated __tensor_flatten__ results and assumed
every element was a tensor now handles opaques: subclass_parametrization,
non_strict_utils, FSDP _init_utils, common_utils (get_untyped_storages),
frontend_utils, and parametrize.py (which stores opaques as plain
attributes rather than parameters).

Tests added (all in test/test_opaque_obj_v2.py):

test_tensor_subclass_with_opaque_attr_backward
test_tensor_subclass_opaque_backward_compiled_autograd
test_tensor_subclass_shared_opaque_remapping
test_shared_opaque_identity_guard
test_shared_direct_opaque_identity_guard
test_tensor_subclass_shared_opaque_backward
test_deeply_nested_tensor_subclass_with_opaque
test_subclass_parametrization_with_opaque_attrs
test_export_non_strict_with_opaque_attrs
test_get_untyped_storages_with_opaque_attrs
test_subclass_opaque_attrs_cache_hit
test_value_type_opaque_in_tensor_attrs_errors

Authored with Claude.

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @kadeng @chauhang @amjames @Lucaskabela @jataylo

The first return value of __tensor_flatten__ can now contain opaque (non-tensor) values such as DeviceMesh. Opaques flow through the flat arg list alongside tensors: they get indices, become graph inputs/outputs, and are pulled from all_args by position during subclass reconstruction — the same as PlainTensorMeta. An empty OpaqueMeta marker in SubclassCreationMeta.attrs distinguishes opaque slots from tensor slots (needed by process_runtime_tangent, which skips non-differentiable opaques). All code that previously iterated __tensor_flatten__ results and assumed every element was a tensor now handles opaques: subclass_parametrization, non_strict_utils, FSDP _init_utils, common_utils (get_untyped_storages), frontend_utils, and parametrize.py (which stores opaques as plain attributes rather than parameters). Tests added (all in test/test_opaque_obj_v2.py): - test_tensor_subclass_with_opaque_attr_backward - test_tensor_subclass_opaque_backward_compiled_autograd - test_tensor_subclass_shared_opaque_remapping - test_shared_opaque_identity_guard - test_shared_direct_opaque_identity_guard - test_tensor_subclass_shared_opaque_backward - test_deeply_nested_tensor_subclass_with_opaque - test_subclass_parametrization_with_opaque_attrs - test_export_non_strict_with_opaque_attrs - test_get_untyped_storages_with_opaque_attrs - test_subclass_opaque_attrs_cache_hit Authored with Claude. [ghstack-poisoned]

pytorch-bot · 2026-03-04T18:13:59Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki

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

✅ You can merge normally! (2 Unrelated Failures)

As of commit f583a4b with merge base 9a8ed4e ():

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

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

trunk / macos-py3-arm64 / test (mps, 1, 1, macos-m1-14) (gh) (trunk failure)
Build left local git repository checkout dirty

UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:

inductor / inductor-cpu-test / test (cpu_inductor_torchbench, 1, 2, linux.2xlarge.amx, unstable) (gh) (#174929)
detectron2_maskrcnn_r_50_fpn

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

The first return value of __tensor_flatten__ can now contain opaque (non-tensor) values such as DeviceMesh. Opaques flow through the flat arg list alongside tensors: they get indices, become graph inputs/outputs, and are pulled from all_args by position during subclass reconstruction — the same as PlainTensorMeta. An empty OpaqueMeta marker in SubclassCreationMeta.attrs distinguishes opaque slots from tensor slots (needed by process_runtime_tangent, which skips non-differentiable opaques). All code that previously iterated __tensor_flatten__ results and assumed every element was a tensor now handles opaques: subclass_parametrization, non_strict_utils, FSDP _init_utils, common_utils (get_untyped_storages), frontend_utils, and parametrize.py (which stores opaques as plain attributes rather than parameters). Tests added (all in test/test_opaque_obj_v2.py): - test_tensor_subclass_with_opaque_attr_backward - test_tensor_subclass_opaque_backward_compiled_autograd - test_tensor_subclass_shared_opaque_remapping - test_shared_opaque_identity_guard - test_shared_direct_opaque_identity_guard - test_tensor_subclass_shared_opaque_backward - test_deeply_nested_tensor_subclass_with_opaque - test_subclass_parametrization_with_opaque_attrs - test_export_non_strict_with_opaque_attrs - test_get_untyped_storages_with_opaque_attrs - test_subclass_opaque_attrs_cache_hit Authored with Claude. ghstack-source-id: f485fef Pull Request resolved: #176457

The first return value of `__tensor_flatten__` can now contain opaque (non-tensor) values such as DeviceMesh. Opaques flow through the flat arg list alongside tensors: they get indices, become graph inputs/outputs, and are pulled from all_args by position during subclass reconstruction — the same as PlainTensorMeta. An empty OpaqueMeta marker in SubclassCreationMeta.attrs distinguishes opaque slots from tensor slots (needed by process_runtime_tangent, which skips non-differentiable opaques). All code that previously iterated `__tensor_flatten__` results and assumed every element was a tensor now handles opaques: subclass_parametrization, non_strict_utils, FSDP _init_utils, common_utils (get_untyped_storages), frontend_utils, and parametrize.py (which stores opaques as plain attributes rather than parameters). Tests added (all in test/test_opaque_obj_v2.py): - test_tensor_subclass_with_opaque_attr_backward - test_tensor_subclass_opaque_backward_compiled_autograd - test_tensor_subclass_shared_opaque_remapping - test_shared_opaque_identity_guard - test_shared_direct_opaque_identity_guard - test_tensor_subclass_shared_opaque_backward - test_deeply_nested_tensor_subclass_with_opaque - test_subclass_parametrization_with_opaque_attrs - test_export_non_strict_with_opaque_attrs - test_get_untyped_storages_with_opaque_attrs - test_subclass_opaque_attrs_cache_hit Authored with Claude. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx kadeng chauhang amjames Lucaskabela jataylo [ghstack-poisoned]

The first return value of __tensor_flatten__ can now contain opaque (non-tensor) values such as DeviceMesh. Opaques flow through the flat arg list alongside tensors: they get indices, become graph inputs/outputs, and are pulled from all_args by position during subclass reconstruction — the same as PlainTensorMeta. An empty OpaqueMeta marker in SubclassCreationMeta.attrs distinguishes opaque slots from tensor slots (needed by process_runtime_tangent, which skips non-differentiable opaques). All code that previously iterated __tensor_flatten__ results and assumed every element was a tensor now handles opaques: subclass_parametrization, non_strict_utils, FSDP _init_utils, common_utils (get_untyped_storages), frontend_utils, and parametrize.py (which stores opaques as plain attributes rather than parameters). Tests added (all in test/test_opaque_obj_v2.py): - test_tensor_subclass_with_opaque_attr_backward - test_tensor_subclass_opaque_backward_compiled_autograd - test_tensor_subclass_shared_opaque_remapping - test_shared_opaque_identity_guard - test_shared_direct_opaque_identity_guard - test_tensor_subclass_shared_opaque_backward - test_deeply_nested_tensor_subclass_with_opaque - test_subclass_parametrization_with_opaque_attrs - test_export_non_strict_with_opaque_attrs - test_get_untyped_storages_with_opaque_attrs - test_subclass_opaque_attrs_cache_hit Authored with Claude. ghstack-source-id: a18076d Pull Request resolved: #176457

The first return value of `__tensor_flatten__` can now contain opaque (non-tensor) values such as DeviceMesh. Opaques flow through the flat arg list alongside tensors: they get indices, become graph inputs/outputs, and are pulled from all_args by position during subclass reconstruction — the same as PlainTensorMeta. An empty OpaqueMeta marker in SubclassCreationMeta.attrs distinguishes opaque slots from tensor slots (needed by process_runtime_tangent, which skips non-differentiable opaques). All code that previously iterated `__tensor_flatten__` results and assumed every element was a tensor now handles opaques: subclass_parametrization, non_strict_utils, FSDP _init_utils, common_utils (get_untyped_storages), frontend_utils, and parametrize.py (which stores opaques as plain attributes rather than parameters). Tests added (all in test/test_opaque_obj_v2.py): - test_tensor_subclass_with_opaque_attr_backward - test_tensor_subclass_opaque_backward_compiled_autograd - test_tensor_subclass_shared_opaque_remapping - test_shared_opaque_identity_guard - test_shared_direct_opaque_identity_guard - test_tensor_subclass_shared_opaque_backward - test_deeply_nested_tensor_subclass_with_opaque - test_subclass_parametrization_with_opaque_attrs - test_export_non_strict_with_opaque_attrs - test_get_untyped_storages_with_opaque_attrs - test_subclass_opaque_attrs_cache_hit - test_value_type_opaque_in_tensor_attrs_errors Authored with Claude. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx kadeng chauhang amjames Lucaskabela jataylo [ghstack-poisoned]

The first return value of __tensor_flatten__ can now contain opaque (non-tensor) values such as DeviceMesh. Opaques flow through the flat arg list alongside tensors: they get indices, become graph inputs/outputs, and are pulled from all_args by position during subclass reconstruction — the same as PlainTensorMeta. An empty OpaqueMeta marker in SubclassCreationMeta.attrs distinguishes opaque slots from tensor slots (needed by process_runtime_tangent, which skips non-differentiable opaques). All code that previously iterated __tensor_flatten__ results and assumed every element was a tensor now handles opaques: subclass_parametrization, non_strict_utils, FSDP _init_utils, common_utils (get_untyped_storages), frontend_utils, and parametrize.py (which stores opaques as plain attributes rather than parameters). Tests added (all in test/test_opaque_obj_v2.py): - test_tensor_subclass_with_opaque_attr_backward - test_tensor_subclass_opaque_backward_compiled_autograd - test_tensor_subclass_shared_opaque_remapping - test_shared_opaque_identity_guard - test_shared_direct_opaque_identity_guard - test_tensor_subclass_shared_opaque_backward - test_deeply_nested_tensor_subclass_with_opaque - test_subclass_parametrization_with_opaque_attrs - test_export_non_strict_with_opaque_attrs - test_get_untyped_storages_with_opaque_attrs - test_subclass_opaque_attrs_cache_hit Authored with Claude. ghstack-source-id: 47dad0f Pull Request resolved: #176457

The first return value of `__tensor_flatten__` can now contain opaque (non-tensor) values such as DeviceMesh. Opaques flow through the flat arg list alongside tensors: they get indices, become graph inputs/outputs, and are pulled from all_args by position during subclass reconstruction — the same as PlainTensorMeta. An empty OpaqueMeta marker in SubclassCreationMeta.attrs distinguishes opaque slots from tensor slots (needed by process_runtime_tangent, which skips non-differentiable opaques). All code that previously iterated `__tensor_flatten__` results and assumed every element was a tensor now handles opaques: subclass_parametrization, non_strict_utils, FSDP _init_utils, common_utils (get_untyped_storages), frontend_utils, and parametrize.py (which stores opaques as plain attributes rather than parameters). Tests added (all in test/test_opaque_obj_v2.py): - test_tensor_subclass_with_opaque_attr_backward - test_tensor_subclass_opaque_backward_compiled_autograd - test_tensor_subclass_shared_opaque_remapping - test_shared_opaque_identity_guard - test_shared_direct_opaque_identity_guard - test_tensor_subclass_shared_opaque_backward - test_deeply_nested_tensor_subclass_with_opaque - test_subclass_parametrization_with_opaque_attrs - test_export_non_strict_with_opaque_attrs - test_get_untyped_storages_with_opaque_attrs - test_subclass_opaque_attrs_cache_hit - test_value_type_opaque_in_tensor_attrs_errors Authored with Claude. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx kadeng chauhang amjames Lucaskabela jataylo [ghstack-poisoned]

The first return value of __tensor_flatten__ can now contain opaque (non-tensor) values such as DeviceMesh. Opaques flow through the flat arg list alongside tensors: they get indices, become graph inputs/outputs, and are pulled from all_args by position during subclass reconstruction — the same as PlainTensorMeta. An empty OpaqueMeta marker in SubclassCreationMeta.attrs distinguishes opaque slots from tensor slots (needed by process_runtime_tangent, which skips non-differentiable opaques). All code that previously iterated __tensor_flatten__ results and assumed every element was a tensor now handles opaques: subclass_parametrization, non_strict_utils, FSDP _init_utils, common_utils (get_untyped_storages), frontend_utils, and parametrize.py (which stores opaques as plain attributes rather than parameters). Tests added (all in test/test_opaque_obj_v2.py): - test_tensor_subclass_with_opaque_attr_backward - test_tensor_subclass_opaque_backward_compiled_autograd - test_tensor_subclass_shared_opaque_remapping - test_shared_opaque_identity_guard - test_shared_direct_opaque_identity_guard - test_tensor_subclass_shared_opaque_backward - test_deeply_nested_tensor_subclass_with_opaque - test_subclass_parametrization_with_opaque_attrs - test_export_non_strict_with_opaque_attrs - test_get_untyped_storages_with_opaque_attrs - test_subclass_opaque_attrs_cache_hit Authored with Claude. ghstack-source-id: 2dd7af8 Pull Request resolved: #176457

The first return value of `__tensor_flatten__` can now contain opaque (non-tensor) values such as DeviceMesh. Opaques flow through the flat arg list alongside tensors: they get indices, become graph inputs/outputs, and are pulled from all_args by position during subclass reconstruction — the same as PlainTensorMeta. An empty OpaqueMeta marker in SubclassCreationMeta.attrs distinguishes opaque slots from tensor slots (needed by process_runtime_tangent, which skips non-differentiable opaques). All code that previously iterated `__tensor_flatten__` results and assumed every element was a tensor now handles opaques: subclass_parametrization, non_strict_utils, FSDP _init_utils, common_utils (get_untyped_storages), frontend_utils, and parametrize.py (which stores opaques as plain attributes rather than parameters). Tests added (all in test/test_opaque_obj_v2.py): - test_tensor_subclass_with_opaque_attr_backward - test_tensor_subclass_opaque_backward_compiled_autograd - test_tensor_subclass_shared_opaque_remapping - test_shared_opaque_identity_guard - test_shared_direct_opaque_identity_guard - test_tensor_subclass_shared_opaque_backward - test_deeply_nested_tensor_subclass_with_opaque - test_subclass_parametrization_with_opaque_attrs - test_export_non_strict_with_opaque_attrs - test_get_untyped_storages_with_opaque_attrs - test_subclass_opaque_attrs_cache_hit - test_value_type_opaque_in_tensor_attrs_errors Authored with Claude. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx kadeng chauhang amjames Lucaskabela jataylo [ghstack-poisoned]

The first return value of __tensor_flatten__ can now contain opaque (non-tensor) values such as DeviceMesh. Opaques flow through the flat arg list alongside tensors: they get indices, become graph inputs/outputs, and are pulled from all_args by position during subclass reconstruction — the same as PlainTensorMeta. An empty OpaqueMeta marker in SubclassCreationMeta.attrs distinguishes opaque slots from tensor slots (needed by process_runtime_tangent, which skips non-differentiable opaques). All code that previously iterated __tensor_flatten__ results and assumed every element was a tensor now handles opaques: subclass_parametrization, non_strict_utils, FSDP _init_utils, common_utils (get_untyped_storages), frontend_utils, and parametrize.py (which stores opaques as plain attributes rather than parameters). Tests added (all in test/test_opaque_obj_v2.py): - test_tensor_subclass_with_opaque_attr_backward - test_tensor_subclass_opaque_backward_compiled_autograd - test_tensor_subclass_shared_opaque_remapping - test_shared_opaque_identity_guard - test_shared_direct_opaque_identity_guard - test_tensor_subclass_shared_opaque_backward - test_deeply_nested_tensor_subclass_with_opaque - test_subclass_parametrization_with_opaque_attrs - test_export_non_strict_with_opaque_attrs - test_get_untyped_storages_with_opaque_attrs - test_subclass_opaque_attrs_cache_hit Authored with Claude. ghstack-source-id: 2dd7af8 Pull Request resolved: #176457

torch/fx/experimental/symbolic_shapes.py

torch/_dynamo/variables/builder.py

torch/_dynamo/variables/tensor.py

The first return value of __tensor_flatten__ can now contain opaque (non-tensor) values such as DeviceMesh. Opaques flow through the flat arg list alongside tensors: they get indices, become graph inputs/outputs, and are pulled from all_args by position during subclass reconstruction — the same as PlainTensorMeta. An empty OpaqueMeta marker in SubclassCreationMeta.attrs distinguishes opaque slots from tensor slots (needed by process_runtime_tangent, which skips non-differentiable opaques). All code that previously iterated __tensor_flatten__ results and assumed every element was a tensor now handles opaques: subclass_parametrization, non_strict_utils, FSDP _init_utils, common_utils (get_untyped_storages), frontend_utils, and parametrize.py (which stores opaques as plain attributes rather than parameters). Tests added (all in test/test_opaque_obj_v2.py): - test_tensor_subclass_with_opaque_attr_backward - test_tensor_subclass_opaque_backward_compiled_autograd - test_tensor_subclass_shared_opaque_remapping - test_shared_opaque_identity_guard - test_shared_direct_opaque_identity_guard - test_tensor_subclass_shared_opaque_backward - test_deeply_nested_tensor_subclass_with_opaque - test_subclass_parametrization_with_opaque_attrs - test_export_non_strict_with_opaque_attrs - test_get_untyped_storages_with_opaque_attrs - test_subclass_opaque_attrs_cache_hit Authored with Claude. ghstack-source-id: bc4d79e Pull Request resolved: #176457

The first return value of `__tensor_flatten__` can now contain opaque (non-tensor) values such as DeviceMesh. Opaques flow through the flat arg list alongside tensors: they get indices, become graph inputs/outputs, and are pulled from all_args by position during subclass reconstruction — the same as PlainTensorMeta. An empty OpaqueMeta marker in SubclassCreationMeta.attrs distinguishes opaque slots from tensor slots (needed by process_runtime_tangent, which skips non-differentiable opaques). All code that previously iterated `__tensor_flatten__` results and assumed every element was a tensor now handles opaques: subclass_parametrization, non_strict_utils, FSDP _init_utils, common_utils (get_untyped_storages), frontend_utils, and parametrize.py (which stores opaques as plain attributes rather than parameters). Tests added (all in test/test_opaque_obj_v2.py): - test_tensor_subclass_with_opaque_attr_backward - test_tensor_subclass_opaque_backward_compiled_autograd - test_tensor_subclass_shared_opaque_remapping - test_shared_opaque_identity_guard - test_shared_direct_opaque_identity_guard - test_tensor_subclass_shared_opaque_backward - test_deeply_nested_tensor_subclass_with_opaque - test_subclass_parametrization_with_opaque_attrs - test_export_non_strict_with_opaque_attrs - test_get_untyped_storages_with_opaque_attrs - test_subclass_opaque_attrs_cache_hit - test_value_type_opaque_in_tensor_attrs_errors Authored with Claude. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx kadeng chauhang amjames Lucaskabela jataylo [ghstack-poisoned]

The first return value of __tensor_flatten__ can now contain opaque (non-tensor) values such as DeviceMesh. Opaques flow through the flat arg list alongside tensors: they get indices, become graph inputs/outputs, and are pulled from all_args by position during subclass reconstruction — the same as PlainTensorMeta. An empty OpaqueMeta marker in SubclassCreationMeta.attrs distinguishes opaque slots from tensor slots (needed by process_runtime_tangent, which skips non-differentiable opaques). All code that previously iterated __tensor_flatten__ results and assumed every element was a tensor now handles opaques: subclass_parametrization, non_strict_utils, FSDP _init_utils, common_utils (get_untyped_storages), frontend_utils, and parametrize.py (which stores opaques as plain attributes rather than parameters). Tests added (all in test/test_opaque_obj_v2.py): - test_tensor_subclass_with_opaque_attr_backward - test_tensor_subclass_opaque_backward_compiled_autograd - test_tensor_subclass_shared_opaque_remapping - test_shared_opaque_identity_guard - test_shared_direct_opaque_identity_guard - test_tensor_subclass_shared_opaque_backward - test_deeply_nested_tensor_subclass_with_opaque - test_subclass_parametrization_with_opaque_attrs - test_export_non_strict_with_opaque_attrs - test_get_untyped_storages_with_opaque_attrs - test_subclass_opaque_attrs_cache_hit Authored with Claude. ghstack-source-id: 65fb56c Pull Request resolved: #176457

The first return value of `__tensor_flatten__` can now contain opaque (non-tensor) values such as DeviceMesh. Opaques flow through the flat arg list alongside tensors: they get indices, become graph inputs/outputs, and are pulled from all_args by position during subclass reconstruction — the same as PlainTensorMeta. An empty OpaqueMeta marker in SubclassCreationMeta.attrs distinguishes opaque slots from tensor slots (needed by process_runtime_tangent, which skips non-differentiable opaques). All code that previously iterated `__tensor_flatten__` results and assumed every element was a tensor now handles opaques: subclass_parametrization, non_strict_utils, FSDP _init_utils, common_utils (get_untyped_storages), frontend_utils, and parametrize.py (which stores opaques as plain attributes rather than parameters). Tests added (all in test/test_opaque_obj_v2.py): - test_tensor_subclass_with_opaque_attr_backward - test_tensor_subclass_opaque_backward_compiled_autograd - test_tensor_subclass_shared_opaque_remapping - test_shared_opaque_identity_guard - test_shared_direct_opaque_identity_guard - test_tensor_subclass_shared_opaque_backward - test_deeply_nested_tensor_subclass_with_opaque - test_subclass_parametrization_with_opaque_attrs - test_export_non_strict_with_opaque_attrs - test_get_untyped_storages_with_opaque_attrs - test_subclass_opaque_attrs_cache_hit - test_value_type_opaque_in_tensor_attrs_errors Authored with Claude. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx kadeng chauhang amjames Lucaskabela jataylo [ghstack-poisoned]

The first return value of __tensor_flatten__ can now contain opaque (non-tensor) values such as DeviceMesh. Opaques flow through the flat arg list alongside tensors: they get indices, become graph inputs/outputs, and are pulled from all_args by position during subclass reconstruction — the same as PlainTensorMeta. An empty OpaqueMeta marker in SubclassCreationMeta.attrs distinguishes opaque slots from tensor slots (needed by process_runtime_tangent, which skips non-differentiable opaques). All code that previously iterated __tensor_flatten__ results and assumed every element was a tensor now handles opaques: subclass_parametrization, non_strict_utils, FSDP _init_utils, common_utils (get_untyped_storages), frontend_utils, and parametrize.py (which stores opaques as plain attributes rather than parameters). Tests added (all in test/test_opaque_obj_v2.py): - test_tensor_subclass_with_opaque_attr_backward - test_tensor_subclass_opaque_backward_compiled_autograd - test_tensor_subclass_shared_opaque_remapping - test_shared_opaque_identity_guard - test_shared_direct_opaque_identity_guard - test_tensor_subclass_shared_opaque_backward - test_deeply_nested_tensor_subclass_with_opaque - test_subclass_parametrization_with_opaque_attrs - test_export_non_strict_with_opaque_attrs - test_get_untyped_storages_with_opaque_attrs - test_subclass_opaque_attrs_cache_hit Authored with Claude. ghstack-source-id: 45b0259 Pull Request resolved: #176457

The first return value of `__tensor_flatten__` can now contain opaque (non-tensor) values such as DeviceMesh. Opaques flow through the flat arg list alongside tensors: they get indices, become graph inputs/outputs, and are pulled from all_args by position during subclass reconstruction — the same as PlainTensorMeta. An empty OpaqueMeta marker in SubclassCreationMeta.attrs distinguishes opaque slots from tensor slots (needed by process_runtime_tangent, which skips non-differentiable opaques). All code that previously iterated `__tensor_flatten__` results and assumed every element was a tensor now handles opaques: subclass_parametrization, non_strict_utils, FSDP _init_utils, common_utils (get_untyped_storages), frontend_utils, and parametrize.py (which stores opaques as plain attributes rather than parameters). Tests added (all in test/test_opaque_obj_v2.py): - test_tensor_subclass_with_opaque_attr_backward - test_tensor_subclass_opaque_backward_compiled_autograd - test_tensor_subclass_shared_opaque_remapping - test_shared_opaque_identity_guard - test_shared_direct_opaque_identity_guard - test_tensor_subclass_shared_opaque_backward - test_deeply_nested_tensor_subclass_with_opaque - test_subclass_parametrization_with_opaque_attrs - test_export_non_strict_with_opaque_attrs - test_get_untyped_storages_with_opaque_attrs - test_subclass_opaque_attrs_cache_hit - test_value_type_opaque_in_tensor_attrs_errors Authored with Claude. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx kadeng chauhang amjames Lucaskabela jataylo [ghstack-poisoned]

The first return value of __tensor_flatten__ can now contain opaque (non-tensor) values such as DeviceMesh. Opaques flow through the flat arg list alongside tensors: they get indices, become graph inputs/outputs, and are pulled from all_args by position during subclass reconstruction — the same as PlainTensorMeta. An empty OpaqueMeta marker in SubclassCreationMeta.attrs distinguishes opaque slots from tensor slots (needed by process_runtime_tangent, which skips non-differentiable opaques). All code that previously iterated __tensor_flatten__ results and assumed every element was a tensor now handles opaques: subclass_parametrization, non_strict_utils, FSDP _init_utils, common_utils (get_untyped_storages), frontend_utils, and parametrize.py (which stores opaques as plain attributes rather than parameters). Tests added (all in test/test_opaque_obj_v2.py): - test_tensor_subclass_with_opaque_attr_backward - test_tensor_subclass_opaque_backward_compiled_autograd - test_tensor_subclass_shared_opaque_remapping - test_shared_opaque_identity_guard - test_shared_direct_opaque_identity_guard - test_tensor_subclass_shared_opaque_backward - test_deeply_nested_tensor_subclass_with_opaque - test_subclass_parametrization_with_opaque_attrs - test_export_non_strict_with_opaque_attrs - test_get_untyped_storages_with_opaque_attrs - test_subclass_opaque_attrs_cache_hit Authored with Claude. ghstack-source-id: c505ee7 Pull Request resolved: #176457

The first return value of `__tensor_flatten__` can now contain opaque (non-tensor) values such as DeviceMesh. Opaques flow through the flat arg list alongside tensors: they get indices, become graph inputs/outputs, and are pulled from all_args by position during subclass reconstruction — the same as PlainTensorMeta. An empty OpaqueMeta marker in SubclassCreationMeta.attrs distinguishes opaque slots from tensor slots (needed by process_runtime_tangent, which skips non-differentiable opaques). All code that previously iterated `__tensor_flatten__` results and assumed every element was a tensor now handles opaques: subclass_parametrization, non_strict_utils, FSDP _init_utils, common_utils (get_untyped_storages), frontend_utils, and parametrize.py (which stores opaques as plain attributes rather than parameters). Tests added (all in test/test_opaque_obj_v2.py): - test_tensor_subclass_with_opaque_attr_backward - test_tensor_subclass_opaque_backward_compiled_autograd - test_tensor_subclass_shared_opaque_remapping - test_shared_opaque_identity_guard - test_shared_direct_opaque_identity_guard - test_tensor_subclass_shared_opaque_backward - test_deeply_nested_tensor_subclass_with_opaque - test_subclass_parametrization_with_opaque_attrs - test_export_non_strict_with_opaque_attrs - test_get_untyped_storages_with_opaque_attrs - test_subclass_opaque_attrs_cache_hit - test_value_type_opaque_in_tensor_attrs_errors Authored with Claude. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx kadeng chauhang amjames Lucaskabela jataylo [ghstack-poisoned]

The first return value of __tensor_flatten__ can now contain opaque (non-tensor) values such as DeviceMesh. Opaques flow through the flat arg list alongside tensors: they get indices, become graph inputs/outputs, and are pulled from all_args by position during subclass reconstruction — the same as PlainTensorMeta. An empty OpaqueMeta marker in SubclassCreationMeta.attrs distinguishes opaque slots from tensor slots (needed by process_runtime_tangent, which skips non-differentiable opaques). All code that previously iterated __tensor_flatten__ results and assumed every element was a tensor now handles opaques: subclass_parametrization, non_strict_utils, FSDP _init_utils, common_utils (get_untyped_storages), frontend_utils, and parametrize.py (which stores opaques as plain attributes rather than parameters). Tests added (all in test/test_opaque_obj_v2.py): - test_tensor_subclass_with_opaque_attr_backward - test_tensor_subclass_opaque_backward_compiled_autograd - test_tensor_subclass_shared_opaque_remapping - test_shared_opaque_identity_guard - test_shared_direct_opaque_identity_guard - test_tensor_subclass_shared_opaque_backward - test_deeply_nested_tensor_subclass_with_opaque - test_subclass_parametrization_with_opaque_attrs - test_export_non_strict_with_opaque_attrs - test_get_untyped_storages_with_opaque_attrs - test_subclass_opaque_attrs_cache_hit Authored with Claude. ghstack-source-id: fa270e1 Pull Request resolved: #176457

aorenste · 2026-03-08T03:56:51Z

@pytorchbot merge

pytorchmergebot · 2026-03-08T03:59:06Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

pytorchmergebot · 2026-03-08T04:20:20Z

Merge failed

Reason: 1 jobs have failed, first few of them are: trunk / macos-py3-arm64 / test (default, 3, 3, macos-m1-stable)

Details for Dev Infra team

Raised by workflow job

The first return value of `__tensor_flatten__` can now contain opaque (non-tensor) values such as DeviceMesh. Opaques flow through the flat arg list alongside tensors: they get indices, become graph inputs/outputs, and are pulled from all_args by position during subclass reconstruction — the same as PlainTensorMeta. An empty OpaqueMeta marker in SubclassCreationMeta.attrs distinguishes opaque slots from tensor slots (needed by process_runtime_tangent, which skips non-differentiable opaques). All code that previously iterated `__tensor_flatten__` results and assumed every element was a tensor now handles opaques: subclass_parametrization, non_strict_utils, FSDP _init_utils, common_utils (get_untyped_storages), frontend_utils, and parametrize.py (which stores opaques as plain attributes rather than parameters). Tests added (all in test/test_opaque_obj_v2.py): - test_tensor_subclass_with_opaque_attr_backward - test_tensor_subclass_opaque_backward_compiled_autograd - test_tensor_subclass_shared_opaque_remapping - test_shared_opaque_identity_guard - test_shared_direct_opaque_identity_guard - test_tensor_subclass_shared_opaque_backward - test_deeply_nested_tensor_subclass_with_opaque - test_subclass_parametrization_with_opaque_attrs - test_export_non_strict_with_opaque_attrs - test_get_untyped_storages_with_opaque_attrs - test_subclass_opaque_attrs_cache_hit - test_value_type_opaque_in_tensor_attrs_errors Authored with Claude. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx kadeng chauhang amjames Lucaskabela jataylo [ghstack-poisoned]

The first return value of __tensor_flatten__ can now contain opaque (non-tensor) values such as DeviceMesh. Opaques flow through the flat arg list alongside tensors: they get indices, become graph inputs/outputs, and are pulled from all_args by position during subclass reconstruction — the same as PlainTensorMeta. An empty OpaqueMeta marker in SubclassCreationMeta.attrs distinguishes opaque slots from tensor slots (needed by process_runtime_tangent, which skips non-differentiable opaques). All code that previously iterated __tensor_flatten__ results and assumed every element was a tensor now handles opaques: subclass_parametrization, non_strict_utils, FSDP _init_utils, common_utils (get_untyped_storages), frontend_utils, and parametrize.py (which stores opaques as plain attributes rather than parameters). Tests added (all in test/test_opaque_obj_v2.py): - test_tensor_subclass_with_opaque_attr_backward - test_tensor_subclass_opaque_backward_compiled_autograd - test_tensor_subclass_shared_opaque_remapping - test_shared_opaque_identity_guard - test_shared_direct_opaque_identity_guard - test_tensor_subclass_shared_opaque_backward - test_deeply_nested_tensor_subclass_with_opaque - test_subclass_parametrization_with_opaque_attrs - test_export_non_strict_with_opaque_attrs - test_get_untyped_storages_with_opaque_attrs - test_subclass_opaque_attrs_cache_hit Authored with Claude. ghstack-source-id: 0dd82ff Pull Request resolved: #176457

The first return value of `__tensor_flatten__` can now contain opaque (non-tensor) values such as DeviceMesh. Opaques flow through the flat arg list alongside tensors: they get indices, become graph inputs/outputs, and are pulled from all_args by position during subclass reconstruction — the same as PlainTensorMeta. An empty OpaqueMeta marker in SubclassCreationMeta.attrs distinguishes opaque slots from tensor slots (needed by process_runtime_tangent, which skips non-differentiable opaques). All code that previously iterated `__tensor_flatten__` results and assumed every element was a tensor now handles opaques: subclass_parametrization, non_strict_utils, FSDP _init_utils, common_utils (get_untyped_storages), frontend_utils, and parametrize.py (which stores opaques as plain attributes rather than parameters). Tests added (all in test/test_opaque_obj_v2.py): - test_tensor_subclass_with_opaque_attr_backward - test_tensor_subclass_opaque_backward_compiled_autograd - test_tensor_subclass_shared_opaque_remapping - test_shared_opaque_identity_guard - test_shared_direct_opaque_identity_guard - test_tensor_subclass_shared_opaque_backward - test_deeply_nested_tensor_subclass_with_opaque - test_subclass_parametrization_with_opaque_attrs - test_export_non_strict_with_opaque_attrs - test_get_untyped_storages_with_opaque_attrs - test_subclass_opaque_attrs_cache_hit - test_value_type_opaque_in_tensor_attrs_errors Authored with Claude. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx kadeng chauhang amjames Lucaskabela jataylo [ghstack-poisoned]

The first return value of __tensor_flatten__ can now contain opaque (non-tensor) values such as DeviceMesh. Opaques flow through the flat arg list alongside tensors: they get indices, become graph inputs/outputs, and are pulled from all_args by position during subclass reconstruction — the same as PlainTensorMeta. An empty OpaqueMeta marker in SubclassCreationMeta.attrs distinguishes opaque slots from tensor slots (needed by process_runtime_tangent, which skips non-differentiable opaques). All code that previously iterated __tensor_flatten__ results and assumed every element was a tensor now handles opaques: subclass_parametrization, non_strict_utils, FSDP _init_utils, common_utils (get_untyped_storages), frontend_utils, and parametrize.py (which stores opaques as plain attributes rather than parameters). Tests added (all in test/test_opaque_obj_v2.py): - test_tensor_subclass_with_opaque_attr_backward - test_tensor_subclass_opaque_backward_compiled_autograd - test_tensor_subclass_shared_opaque_remapping - test_shared_opaque_identity_guard - test_shared_direct_opaque_identity_guard - test_tensor_subclass_shared_opaque_backward - test_deeply_nested_tensor_subclass_with_opaque - test_subclass_parametrization_with_opaque_attrs - test_export_non_strict_with_opaque_attrs - test_get_untyped_storages_with_opaque_attrs - test_subclass_opaque_attrs_cache_hit Authored with Claude. ghstack-source-id: 065031f Pull Request resolved: #176457

The first return value of `__tensor_flatten__` can now contain opaque (non-tensor) values such as DeviceMesh. Opaques flow through the flat arg list alongside tensors: they get indices, become graph inputs/outputs, and are pulled from all_args by position during subclass reconstruction — the same as PlainTensorMeta. An empty OpaqueMeta marker in SubclassCreationMeta.attrs distinguishes opaque slots from tensor slots (needed by process_runtime_tangent, which skips non-differentiable opaques). All code that previously iterated `__tensor_flatten__` results and assumed every element was a tensor now handles opaques: subclass_parametrization, non_strict_utils, FSDP _init_utils, common_utils (get_untyped_storages), frontend_utils, and parametrize.py (which stores opaques as plain attributes rather than parameters). Tests added (all in test/test_opaque_obj_v2.py): - test_tensor_subclass_with_opaque_attr_backward - test_tensor_subclass_opaque_backward_compiled_autograd - test_tensor_subclass_shared_opaque_remapping - test_shared_opaque_identity_guard - test_shared_direct_opaque_identity_guard - test_tensor_subclass_shared_opaque_backward - test_deeply_nested_tensor_subclass_with_opaque - test_subclass_parametrization_with_opaque_attrs - test_export_non_strict_with_opaque_attrs - test_get_untyped_storages_with_opaque_attrs - test_subclass_opaque_attrs_cache_hit - test_value_type_opaque_in_tensor_attrs_errors Authored with Claude. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx kadeng chauhang amjames Lucaskabela jataylo [ghstack-poisoned]

The first return value of __tensor_flatten__ can now contain opaque (non-tensor) values such as DeviceMesh. Opaques flow through the flat arg list alongside tensors: they get indices, become graph inputs/outputs, and are pulled from all_args by position during subclass reconstruction — the same as PlainTensorMeta. An empty OpaqueMeta marker in SubclassCreationMeta.attrs distinguishes opaque slots from tensor slots (needed by process_runtime_tangent, which skips non-differentiable opaques). All code that previously iterated __tensor_flatten__ results and assumed every element was a tensor now handles opaques: subclass_parametrization, non_strict_utils, FSDP _init_utils, common_utils (get_untyped_storages), frontend_utils, and parametrize.py (which stores opaques as plain attributes rather than parameters). Tests added (all in test/test_opaque_obj_v2.py): - test_tensor_subclass_with_opaque_attr_backward - test_tensor_subclass_opaque_backward_compiled_autograd - test_tensor_subclass_shared_opaque_remapping - test_shared_opaque_identity_guard - test_shared_direct_opaque_identity_guard - test_tensor_subclass_shared_opaque_backward - test_deeply_nested_tensor_subclass_with_opaque - test_subclass_parametrization_with_opaque_attrs - test_export_non_strict_with_opaque_attrs - test_get_untyped_storages_with_opaque_attrs - test_subclass_opaque_attrs_cache_hit Authored with Claude. ghstack-source-id: 35ec966 Pull Request resolved: #176457

aorenste · 2026-03-09T02:07:27Z

@pytorchbot merge

pytorchmergebot · 2026-03-09T02:09:37Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

The first return value of `__tensor_flatten__` can now contain opaque (non-tensor) values such as DeviceMesh. Opaques flow through the flat arg list alongside tensors: they get indices, become graph inputs/outputs, and are pulled from all_args by position during subclass reconstruction — the same as PlainTensorMeta. An empty OpaqueMeta marker in SubclassCreationMeta.attrs distinguishes opaque slots from tensor slots (needed by process_runtime_tangent, which skips non-differentiable opaques). All code that previously iterated `__tensor_flatten__` results and assumed every element was a tensor now handles opaques: subclass_parametrization, non_strict_utils, FSDP _init_utils, common_utils (get_untyped_storages), frontend_utils, and parametrize.py (which stores opaques as plain attributes rather than parameters). Tests added (all in test/test_opaque_obj_v2.py): - test_tensor_subclass_with_opaque_attr_backward - test_tensor_subclass_opaque_backward_compiled_autograd - test_tensor_subclass_shared_opaque_remapping - test_shared_opaque_identity_guard - test_shared_direct_opaque_identity_guard - test_tensor_subclass_shared_opaque_backward - test_deeply_nested_tensor_subclass_with_opaque - test_subclass_parametrization_with_opaque_attrs - test_export_non_strict_with_opaque_attrs - test_get_untyped_storages_with_opaque_attrs - test_subclass_opaque_attrs_cache_hit - test_value_type_opaque_in_tensor_attrs_errors Authored with Claude. Pull Request resolved: pytorch#176457 Approved by: https://github.com/ezyang

pytorch-bot bot added ciflow/inductor module: dynamo release notes: export labels Mar 4, 2026

aorenste marked this pull request as ready for review March 5, 2026 19:39

aorenste requested review from angelayi, avikchaudhuri, bdhirsh, bobrenjc93, laithsakka, lezcano, tugsbayasgalan, ydwu4 and zhxchen17 as code owners March 5, 2026 19:39

aorenste requested review from angelayi and ezyang and removed request for angelayi March 5, 2026 19:39

ezyang reviewed Mar 6, 2026

View reviewed changes

torch/fx/experimental/symbolic_shapes.py Outdated Show resolved Hide resolved

ezyang reviewed Mar 6, 2026

View reviewed changes

torch/_dynamo/variables/builder.py Outdated Show resolved Hide resolved

ezyang reviewed Mar 6, 2026

View reviewed changes

torch/_dynamo/variables/tensor.py Outdated Show resolved Hide resolved

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Mar 8, 2026

pytorchmergebot added the merging label Mar 8, 2026

pytorchmergebot removed the merging label Mar 8, 2026

pytorchmergebot added the merging label Mar 9, 2026

pytorchmergebot closed this in e7c0559 Mar 9, 2026

pytorchmergebot added Merged and removed merging labels Mar 9, 2026

github-actions bot deleted the gh/aorenste/206/head branch April 8, 2026 02:24

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support non-tensor attrs in __tensor_flatten__#176457

Support non-tensor attrs in __tensor_flatten__#176457
aorenste wants to merge 13 commits intogh/aorenste/206/basefrom
gh/aorenste/206/head

aorenste commented Mar 4, 2026 •

edited

Loading

Uh oh!

pytorch-bot bot commented Mar 4, 2026 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

aorenste commented Mar 8, 2026

Uh oh!

pytorchmergebot commented Mar 8, 2026

Uh oh!

pytorchmergebot commented Mar 8, 2026

Uh oh!

aorenste commented Mar 9, 2026

Uh oh!

pytorchmergebot commented Mar 9, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Conversation

aorenste commented Mar 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Mar 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

✅ You can merge normally! (2 Unrelated Failures)

Uh oh!

Uh oh!

Uh oh!

Uh oh!

aorenste commented Mar 8, 2026

Uh oh!

pytorchmergebot commented Mar 8, 2026

Merge started

Uh oh!

pytorchmergebot commented Mar 8, 2026

Merge failed

Uh oh!

aorenste commented Mar 9, 2026

Uh oh!

pytorchmergebot commented Mar 9, 2026

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

aorenste commented Mar 4, 2026 •

edited

Loading

pytorch-bot bot commented Mar 4, 2026 •

edited

Loading