Adding view and reduction tags by AlonSardas · Pull Request #153342 · pytorch/pytorch

AlonSardas · 2025-05-11T12:05:17Z

Here's lists of the operator names annotated with view and reduction tags:
view overloads:
as_strided, detach, view_as_real, view_as_complex, diagonal, select, slice, transpose, split, t, expand, view, unsqueeze, unfold, squeeze, permute, unbind, split_with_sizes, alias

reduction overloads:
sum, mean, amin, amax, argmin, argmax, prod, all, norm, var, std, aminmax, nansum, logsumexp, any, std_mean, var_mean, count_nonzero, linalg_vector_norm, max, min

pytorch-bot · 2025-05-11T12:05:20Z

🔗 Helpful Links

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

📄 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 or our office hours

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

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

eellison

nice!

eellison · 2025-05-15T17:24:06Z

test/test_ops.py

Looks great ! Would it be possible to add a test along the lines of https://github.com/pytorch/pytorch/pull/90029/files?show-viewed-files=true&file-filters%5B%5D=#diff-d183f2afc51d6a59bc70094e8f476d2468c45e415500f6eb60abad955e065156R155-R237? wdyt ?

Yes, I added similar tests for reduction and view.
For reduction, the test checks that the output shape is reduced as expected (I needed the reduced_shape function for that).
For view, the test checks that the output tensor is indeed a view of the base tensor.
I also added a similar test for views on regular tensors (not in fake mode), because I didn't see such a test in other places.
What do you think?
There just seems to be a large overhead for these tests (about 1min each) since they iterate over all the operators and not just the relevant view/reduction ops.

I renamed reduced_shape->compute_reduced_shape since it conflicted with variable name inside test_reductions.py. I also added missing documentation for treat_empty_dim_as_none argument

eellison

Looks great ! Thank you for doing this ! Maybe lets get one more pytorch review before landing. cc @bdhirsh, @zou3519 mind weighing in ?

bdhirsh · 2025-05-28T16:45:03Z

aten/src/ATen/native/tags.yaml

+- tag: view
+  desc: |
+          This tag indicates that the operator creates a pure view/alias Tensor and has an explicit
+          derivative formula in derivates.yaml.


hmm view tag is probably less critical / a bit redundant, given that OpOverload objects already tell you whether they are a view based off of their schema (you can run op_overload.is_view, see link)

I guess the counterargument is that tags are more a general-purpose of grouping ops into categories, and it's easier for a pass writer to use tags for everything rather than tags for some metadata and schema info for other metadata. So I'm not too opinionated, although we should at least ensure that our OpInfo tests assert that the tag always matches the op_overload.is_view "source of truth". cc @eellison

I think there should be one source of truth for views. That source of truth should be the schema, not tags. People are not going to add the tags to custom operators.

I'm still concerned about operators whose schemas imply they are views but don't really behave like pure views. This includes reshape, contiguous, resolve_conj, resolve_neg, copy, lift_fresh (and maybe a few more).

Perhaps we should change their schemas to reflect that they're not pure view operators?

bdhirsh · 2025-05-28T16:47:11Z

@eellison it would also be nice to clean up this big hardcoded list of ops in the min-cut partitioner around what is recompute-friendly. I think it is basically a collective of pointwise/reduction/view ops, although someone would probably need to go through the list carefully to ensure that if we port it over to tags we don't change behavior: https://github.com/pytorch/pytorch/blob/main/torch/_functorch/partitioners.py#L1910

eellison · 2025-05-28T16:48:20Z

@bdhirsh - agreed, we could probably get rid of a lot of those, given that we have added

pytorch/torch/_functorch/partitioners.py

Line 2040 in 3b38989

default_recomputable_ops += pointwise_ops()

.

zou3519

I'm not convinced we need a view tag. There should be a single source of truth for if an operator is a view, and that is currently the schema

bdhirsh · 2025-05-28T16:49:33Z

test/test_ops.py

+        ]  # Fix linalg.vector_norm
+
+        # Additional reduction operators
+        reduction_op_names += [


nit: instead of all of the string parsing here, can we use OpOverload objects directly? That way you can just check func in reduction_ops instead of having to parse the names:

reduction_op_names += [ torch.ops.aten.norm.ScalarOpt_dtype, torch.ops.aten.norm.Scalar, ]

The downside is that you need to manually write out each overload too, although I'd argue this is good since it makes things more explicit (you needed to manually ad the tag for each OpOverload anyway)

bdhirsh · 2025-05-28T16:50:31Z

test/test_ops.py

+    def test_view_tag_coverage(self):
+        # These operators have the inferred property is_view according to their declaration in native_functions.yaml
+        # but they are not pure view operators since they create a copy under certain conditions
+        not_view_operators = ["to", "copy", "lift_fresh"]


Why is "copy" in here? https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/native_functions.yaml#L1785

There is the overload copy.t with copy.t.is_view==True, which is registered in register_prim_ops.cpp

bdhirsh · 2025-05-28T16:52:20Z

test/test_ops.py

+        manually_registered_overloads = [
+            "select.t",
+            "numpy_T.a",
+            "split.default",


I did a quick audit, and I don't see split.default in that file? I do see a split.str though: https://github.com/pytorch/pytorch/blob/main/torch/csrc/jit/runtime/register_prim_ops.cpp#L1778

split.default does have alias annotations. Should this test have failed?

Mentioned above, but I think a nice addition to the test would be to assert that op.is_view == Tag.view in op.tags for every op that we have OpInfo testing for!

I think that split.default is registered in

pytorch/torch/csrc/jit/runtime/register_special_ops.cpp

Line 247 in 4d57644

"aten::split(Tensor(a -> *) self, int[] split_sizes, int dim=0) -> Tensor(a)[]"),

During the discussion we said we want to target only the view operators that are CompositeExplicitAutograd Expand Tag Set: views & reductions #129020 (comment)
This is why the test asserts that op.is_view == Tag.view in op.tags but also not overload.has_kernel_for_dispatch_key(DispatchKey.CompositeImplicitAutograd)

eellison

Could we land this with just the reduction tag ? I have a use case for it now.

github-actions · 2025-09-22T15:36:04Z

Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as Stale.
Feel free to remove the Stale label if you feel this was a mistake.
If you are unable to remove the Stale label please contact a maintainer in order to do so.
If you want the bot to never mark this PR stale again, add the no-stale label.
Stale pull requests will automatically be closed after 30 days of inactivity.

eellison · 2025-09-29T15:42:44Z

@AlonSardas mind rebasing with just the reduction changes ? otherwise, i can push to the pr.

AlonSardas · 2025-10-02T14:09:13Z

I'm currently busy and won't be able to rebase it. Feel free to use just the reduction changes.

Add a new 'reduction' tag to tags.yaml and apply it to 98 reduction operator variants across 21 operator families (sum, mean, min, max, argmin, argmax, amin, amax, aminmax, prod, all, any, norm, var, std, std_mean, var_mean, nansum, logsumexp, count_nonzero, linalg_vector_norm). This tag categorizes operators that perform reduction operations, computing aggregate values across one or more dimensions of input tensor(s). Based on PR #153342 ghstack-source-id: 6c9c42f Pull Request resolved: #165146

Add a new 'reduction' tag to tags.yaml and apply it to 98 reduction operator variants across 21 operator families (sum, mean, min, max, argmin, argmax, amin, amax, aminmax, prod, all, any, norm, var, std, std_mean, var_mean, nansum, logsumexp, count_nonzero, linalg_vector_norm). This tag categorizes operators that perform reduction operations, computing aggregate values across one or more dimensions of input tensor(s). This categorization can be useful for analysis, optimization, and compilation tasks. Note: Only dimensional reduction variants (e.g., min.dim, max.dim) are tagged. Simple unary aggregations (min(Tensor), max(Tensor)) and binary/ elementwise operations (min.other, max.other) are excluded. Based on PR #153342 ghstack-source-id: c0e25de Pull Request resolved: #165155

Add a new 'reduction' tag to tags.yaml and apply it to 98 reduction operator variants across 21 operator families (sum, mean, min, max, argmin, argmax, amin, amax, aminmax, prod, all, any, norm, var, std, std_mean, var_mean, nansum, logsumexp, count_nonzero, linalg_vector_norm). This tag categorizes operators that perform reduction operations, computing aggregate values across one or more dimensions of input tensor(s). This categorization can be useful for analysis, optimization, and compilation tasks. Note: Only dimensional reduction variants (e.g., min.dim, max.dim) are tagged. Simple unary aggregations (min(Tensor), max(Tensor)) and binary/ elementwise operations (min.other, max.other) are excluded. Based on PR #153342 ghstack-source-id: cbb7930 Pull Request resolved: #165155

Add a new 'reduction' tag to tags.yaml and apply it to 98 reduction operator variants across 21 operator families (sum, mean, min, max, argmin, argmax, amin, amax, aminmax, prod, all, any, norm, var, std, std_mean, var_mean, nansum, logsumexp, count_nonzero, linalg_vector_norm). This tag categorizes operators that perform reduction operations, computing aggregate values across one or more dimensions of input tensor(s). This categorization can be useful for analysis, optimization, and compilation tasks. Note: Only dimensional reduction variants (e.g., min.dim, max.dim) are tagged. Simple unary aggregations (min(Tensor), max(Tensor)) and binary/ elementwise operations (min.other, max.other) are excluded. Based on PR #153342 ghstack-source-id: 2f58603 Pull Request resolved: #165155

@AlonSardas

Add a new 'reduction' tag to tags.yaml and apply it to 98 reduction operator variants across 21 operator families (sum, mean, min, max, argmin, argmax, amin, amax, aminmax, prod, all, any, norm, var, std, std_mean, var_mean, nansum, logsumexp, count_nonzero, linalg_vector_norm). This tag categorizes operators that perform reduction operations, computing aggregate values across one or more dimensions of input tensor(s). Based on PR #153342 - co-written with @AlonSardas. Just as we have pointwise tag - this can be useful for compiler passes, or for opting into sharding rules. Pull Request resolved: #165155 Approved by: https://github.com/ezyang, https://github.com/zou3519, https://github.com/mlazos

@AlonSardas

Add a new 'reduction' tag to tags.yaml and apply it to 98 reduction operator variants across 21 operator families (sum, mean, min, max, argmin, argmax, amin, amax, aminmax, prod, all, any, norm, var, std, std_mean, var_mean, nansum, logsumexp, count_nonzero, linalg_vector_norm). This tag categorizes operators that perform reduction operations, computing aggregate values across one or more dimensions of input tensor(s). Based on PR pytorch#153342 - co-written with @AlonSardas. Just as we have pointwise tag - this can be useful for compiler passes, or for opting into sharding rules. Pull Request resolved: pytorch#165155 Approved by: https://github.com/ezyang, https://github.com/zou3519, https://github.com/mlazos

@AlonSardas

Add a new 'reduction' tag to tags.yaml and apply it to 98 reduction operator variants across 21 operator families (sum, mean, min, max, argmin, argmax, amin, amax, aminmax, prod, all, any, norm, var, std, std_mean, var_mean, nansum, logsumexp, count_nonzero, linalg_vector_norm). This tag categorizes operators that perform reduction operations, computing aggregate values across one or more dimensions of input tensor(s). Based on PR pytorch#153342 - co-written with @AlonSardas. Just as we have pointwise tag - this can be useful for compiler passes, or for opting into sharding rules. Pull Request resolved: pytorch#165155 Approved by: https://github.com/ezyang, https://github.com/zou3519, https://github.com/mlazos

AlonSardas requested review from ezyang and mruberry as code owners May 11, 2025 12:05

pytorchbot added the open source label May 11, 2025

eellison requested review from bdhirsh and eellison May 15, 2025 17:15

eellison reviewed May 15, 2025

View reviewed changes

AlonSardas force-pushed the pr129020-reduction-view-tags branch from 75477ec to 9c14c55 Compare May 20, 2025 15:12

AlonSardas added 2 commits May 21, 2025 21:21

Adding view and reduction tags

3de44e5

Adding tests for view and reduction operators with fake tensors

0d5fb10

AlonSardas force-pushed the pr129020-reduction-view-tags branch from 9c14c55 to 0d5fb10 Compare May 21, 2025 18:40

eellison approved these changes May 21, 2025

View reviewed changes

eellison requested a review from zou3519 May 21, 2025 21:40

bdhirsh reviewed May 28, 2025

View reviewed changes

zou3519 requested changes May 28, 2025

View reviewed changes

bdhirsh reviewed May 28, 2025

View reviewed changes

eellison reviewed Jul 24, 2025

View reviewed changes

github-actions bot added the Stale label Sep 22, 2025

eellison mentioned this pull request Sep 29, 2025

Remove decomposition from softmax meta-pytorch/autoparallel#171

Merged

eellison mentioned this pull request Oct 10, 2025

[ATen] Add reduction tag to reduction operators #165155

Closed

eellison mentioned this pull request Oct 10, 2025

[ATen] Add reduction tag to reduction operators #165146

Closed

github-actions bot closed this Nov 1, 2025

CloseChoice mentioned this pull request Dec 9, 2025

Add reduction tags #169944

Closed

Conversation

AlonSardas commented May 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented May 11, 2025 • 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/153342

Uh oh!

eellison left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

eellison left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

AlonSardas Jun 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

bdhirsh commented May 28, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

eellison commented May 28, 2025

Uh oh!

zou3519 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

AlonSardas Jun 4, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

eellison left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Sep 22, 2025

Uh oh!

eellison commented Sep 29, 2025

Uh oh!

AlonSardas commented Oct 2, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

AlonSardas commented May 11, 2025 •

edited

Loading

pytorch-bot bot commented May 11, 2025 •

edited

Loading

eellison left a comment •

edited

Loading

AlonSardas Jun 4, 2025 •

edited

Loading

bdhirsh commented May 28, 2025 •

edited

Loading

AlonSardas Jun 4, 2025 •

edited

Loading