autograd: Add VJP and JVP rules for aten::aminmax #158241

vijayabhaskar-ev · 2025-07-14T15:59:12Z

Adds functionally correct backward (VJP) and forward (JVP) autograd rules for the aten::aminmax operator to derivatives.yaml using existing helper functions. This ensures correct eager mode differentiation.

Fixes #148808

pytorch-bot · 2025-07-14T15:59:16Z

🔗 Helpful Links

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

📄 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.

❌ 1 New Failure

As of commit b7818f6 with merge base c8c221c ():

NEW FAILURE - The following job has failed:

Check Labels / Check labels (gh)
RuntimeError: Error checking labels: PR does not have required labels

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

github-actions · 2025-07-14T15:59:45Z

This PR needs a `release notes:` label

If your changes are user facing and intended to be a part of release notes, please use a label starting with release notes:.

If not, please add the topic: not user facing label.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "topic: not user facing"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

Adds functionally correct backward (VJP) and forward (JVP) autograd rules for the aten::aminmax operator to derivatives.yaml using existing helper functions. This ensures correct eager mode differentiation. Fixes pytorch#148808

… to handle amin and amax correctly. - Modified derivatives.yaml to reflect the changes in autograd behavior. - Updated test cases in test_autograd.py to validate the updated rules. - Updated common_methods_invocations.py to include amin and amax in relevant test cases.

Makes backward pass more efficient by: - Avoiding zeros tensor creation when only one gradient defined - Using in-place operations where possible

• Added a testcase to validate the new logic.

…inmax

…ed testcase

…abhaskar-ev/pytorch into fix-aminmax-autograd-rules

vijayabhaskar-ev · 2025-07-16T05:11:57Z

@soulitzer Adding restore_reduced_dims on grad_max and grad_min is causing Segmentation fault (core dumped)

…kward

vijayabhaskar-ev · 2025-07-16T16:00:46Z

@soulitzer Adding restore_reduced_dims on grad_max and grad_min is causing Segmentation fault (core dumped)

Removed restore_reduced_dims on grad_min/grad_max in aminmax_backward.
scale_grad_by_count already handles broadcasting and the extra expansion created zero stride views that caused gradcheck/functorch crashes. Reverting to the earlier logic fixes the seg fault without altering gradient results.

soulitzer · 2025-07-17T13:21:09Z

scale_grad_by_count already handles broadcasting

Broadcasting doesn't work when certain dimensions are missing. restore_reduced_dims is responsible for restoring those dimensions so broadcasting can happen.

vijayabhaskar-ev · 2025-07-17T13:40:26Z

scale_grad_by_count already handles broadcasting

Broadcasting doesn't work when certain dimensions are missing. restore_reduced_dims is responsible for restoring those dimensions so broadcasting can happen.

If I add restore_reduced_dims on grad_max and grad_min, it's causing Segmentation fault (core dumped).
auto max_reduced = restore_reduced_dims(max, dims, keepdim); auto max_mask = (self == max_reduced); auto grad_max_expanded = restore_reduced_dims(grad_max, dims, keepdim); auto grad_max_result = scale_grad_by_count(grad_max_expanded, max_mask, dims);

 This is the commit that caused the seg fault : https://github.com/pytorch/pytorch/pull/158241/commits/c1ad91c9244a691b5e574f04724587574a64f52a

vijayabhaskar-ev · 2025-07-17T13:44:05Z

scale_grad_by_count already handles broadcasting

Broadcasting doesn't work when certain dimensions are missing. restore_reduced_dims is responsible for restoring those dimensions so broadcasting can happen.

If I add restore_reduced_dims on grad_max and grad_min, it's causing Segmentation fault (core dumped). auto max_reduced = restore_reduced_dims(max, dims, keepdim); auto max_mask = (self == max_reduced); auto grad_max_expanded = restore_reduced_dims(grad_max, dims, keepdim); auto grad_max_result = scale_grad_by_count(grad_max_expanded, max_mask, dims);
 This is the commit that caused the seg fault : https://github.com/pytorch/pytorch/pull/158241/commits/c1ad91c9244a691b5e574f04724587574a64f52a

Let me rebuild and try again

…ad_max

vijayabhaskar-ev · 2025-07-19T17:44:50Z

scale_grad_by_count already handles broadcasting

Broadcasting doesn't work when certain dimensions are missing. restore_reduced_dims is responsible for restoring those dimensions so broadcasting can happen.

If I add restore_reduced_dims on grad_max and grad_min, it's causing Segmentation fault (core dumped). auto max_reduced = restore_reduced_dims(max, dims, keepdim); auto max_mask = (self == max_reduced); auto grad_max_expanded = restore_reduced_dims(grad_max, dims, keepdim); auto grad_max_result = scale_grad_by_count(grad_max_expanded, max_mask, dims);
 This is the commit that caused the seg fault : https://github.com/pytorch/pytorch/pull/158241/commits/c1ad91c9244a691b5e574f04724587574a64f52a
Let me rebuild and try again

@soulitzer PR updated.

soulitzer · 2025-07-21T14:52:32Z

torch/csrc/autograd/FunctionsManual.cpp

+    auto grad_min_full =
+        restore_reduced_dims(grad_min, dims, keepdim)
+            .expand_as(min_mask)          
+            .contiguous();


this doesn't look right, we don't want a copy

Removed contiguous() as it cause's extra memory footprint

vijayabhaskar-ev · 2025-08-11T07:48:14Z

@soulitzer Can you review this PR? Resolved the issue's mentioned in the comments.

vijayabhaskar-ev requested review from mruberry, albanD and soulitzer as code owners July 14, 2025 15:59

pytorchbot added the open source label Jul 14, 2025

albanD removed their request for review July 14, 2025 18:16

albanD added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Jul 14, 2025

vijayabhaskar-ev added 8 commits July 15, 2025 12:09

autograd: Add VJP and JVP rules for aten::aminmax

de8d968

Adds functionally correct backward (VJP) and forward (JVP) autograd rules for the aten::aminmax operator to derivatives.yaml using existing helper functions. This ensures correct eager mode differentiation. Fixes pytorch#148808

[autograd] Optimize tensor allocation in aminmax_backward

27a560a

Makes backward pass more efficient by: - Avoiding zeros tensor creation when only one gradient defined - Using in-place operations where possible

• Modified the mask generation logic.

878649a

• Added a testcase to validate the new logic.

Refactored aminmax_backward and added testcases for amin, amax and am…

4285acb

…inmax

Update dimension parameter to use std::optional<int64_t> and refactor…

9142dbc

…ed testcase

Merge branch 'fix-aminmax-autograd-rules' of https://github.com/vijay…

93d915e

…abhaskar-ev/pytorch into fix-aminmax-autograd-rules

Added restore_reduced_dims for grad_max and grad_min

c1ad91c

Removed restore_reduced_dims fro grad_min and grad_max in aminmax_bac…

889bbf3

…kward

vijayabhaskar-ev force-pushed the fix-aminmax-autograd-rules branch from 4c871fc to 889bbf3 Compare July 16, 2025 15:57

Refactored aminmax_backward and added reduced_dim for grad_min and gr…

610b7ea

…ad_max

soulitzer reviewed Jul 21, 2025

View reviewed changes

vijayabhaskar-ev added 2 commits July 21, 2025 20:39

Removed contiguous() call in aminmax_backward

f7e09b0

Removed expand_as() call in aminmax_backward

b7818f6

vijayabhaskar-ev requested a review from soulitzer July 29, 2025 10:04

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

autograd: Add VJP and JVP rules for aten::aminmax #158241

autograd: Add VJP and JVP rules for aten::aminmax #158241

vijayabhaskar-ev commented Jul 14, 2025 •

edited

Loading

Uh oh!

pytorch-bot bot commented Jul 14, 2025 •

edited

Loading

Uh oh!

github-actions bot commented Jul 14, 2025

Uh oh!

vijayabhaskar-ev commented Jul 16, 2025

Uh oh!

vijayabhaskar-ev commented Jul 16, 2025

Uh oh!

soulitzer commented Jul 17, 2025

Uh oh!

vijayabhaskar-ev commented Jul 17, 2025

Uh oh!

vijayabhaskar-ev commented Jul 17, 2025

Uh oh!

vijayabhaskar-ev commented Jul 19, 2025

Uh oh!

soulitzer Jul 21, 2025

Uh oh!

vijayabhaskar-ev Jul 21, 2025

Uh oh!

vijayabhaskar-ev commented Aug 11, 2025

Uh oh!

Uh oh!

autograd: Add VJP and JVP rules for aten::aminmax #158241

Are you sure you want to change the base?

autograd: Add VJP and JVP rules for aten::aminmax #158241

Conversation

vijayabhaskar-ev commented Jul 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Jul 14, 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/158241

❌ 1 New Failure

Uh oh!

github-actions bot commented Jul 14, 2025

This PR needs a release notes: label

Uh oh!

vijayabhaskar-ev commented Jul 16, 2025

Uh oh!

vijayabhaskar-ev commented Jul 16, 2025

Uh oh!

soulitzer commented Jul 17, 2025

Uh oh!

vijayabhaskar-ev commented Jul 17, 2025

Uh oh!

vijayabhaskar-ev commented Jul 17, 2025

Uh oh!

vijayabhaskar-ev commented Jul 19, 2025

Uh oh!

soulitzer Jul 21, 2025

Choose a reason for hiding this comment

Uh oh!

vijayabhaskar-ev Jul 21, 2025

Choose a reason for hiding this comment

Uh oh!

vijayabhaskar-ev commented Aug 11, 2025

Uh oh!

Uh oh!

vijayabhaskar-ev commented Jul 14, 2025 •

edited

Loading

pytorch-bot bot commented Jul 14, 2025 •

edited

Loading

This PR needs a `release notes:` label