[dtensor] Add single_dim_strategy functions for reduction-with-indices ops by anshul-si · Pull Request #179200 · pytorch/pytorch

anshul-si · 2026-04-02T23:34:00Z

Stack from ghstack (oldest at bottom):

[dtensor][math_ops] strategies for remaining ops #179203
[dtensor] Add single_dim_strategy functions for pooling, replicate-only, and histc ops #179202
[dtensor] Add single_dim_strategy functions for reduction ops #179201
-> [dtensor] Add single_dim_strategy functions for reduction-with-indices ops #179200
[dtensor][math_ops] Add single_dim_strategy functions for sort-like, scan and softmax ops #179068

Refactor max.dim, min.dim, median.dim, nanmedian.dim, mode, kthvalue to
use the shared _shard_non_reduction_dim helper. Add new single_dim_strategy
for argmax/argmin (decorator commented out, old register_op_strategy kept).

…s ops Refactor max.dim, min.dim, median.dim, nanmedian.dim, mode, kthvalue to use the shared _shard_non_reduction_dim helper. Add new single_dim_strategy for argmax/argmin (decorator commented out, old register_op_strategy kept). [ghstack-poisoned]

pytorch-bot · 2026-04-02T23:34:05Z

🔗 Helpful Links

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

📄 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! (5 Unrelated Failures)

As of commit a754411 with merge base c54d05a ():

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

pull / linux-jammy-py3.10-gcc11 / test-osdc (distributed, 2, 3, mt-l-x86iamx-8-64) (gh) (disabled by #187020)
test/distributed/tensor/test_op_strategy.py::TestCostModel::test_redistribute_cost_latency
torchtitan-test / torchtitan-x-pytorch-test / test-osdc (torchtitan_models_integration, 1, 1, mt-l-x86aavx2-189-704-a10g-8) (gh) (similar failure)
[OSDC] Step script exited with code 1. This is a script/workflow error, not an infrastructure issue. Check the step logs above for the actual failure.

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

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

torchtitan-test / torchtitan-x-pytorch-test / test-osdc (torchtitan_features_integration, 1, 1, mt-l-x86aavx2-189-704-a10g-8) (gh) (trunk failure)
RuntimeError: 1 test steps failed: ['scripts/ci/pytorch_ci_test_runner.sh feature_tests']
trunk / linux-jammy-rocm-py3.10-mi355 / test (default, 7, 8, linux.rocm.gpu.gfx950.1) (gh) (trunk failure)
test/test_cuda.py::TestCuda::test_out_of_memory_retry
trunk / macos-py3-arm64 / test (mps, 1, 1, macos-m2-15) (gh) (trunk failure)
Process completed with exit code 139.

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

…with-indices ops" Refactor max.dim, min.dim, median.dim, nanmedian.dim, mode, kthvalue to use the shared _shard_non_reduction_dim helper. Add new single_dim_strategy for argmax/argmin (decorator commented out, old register_op_strategy kept). [ghstack-poisoned]

zpcore · 2026-04-14T06:48:26Z

+# @register_single_dim_strategy(
+#     list(ARGMAX_ARGMIN_OPS.keys()),
+#     schema_info=RuntimeSchemaInfo(1),
+#     allow_uneven_sharding=True,
+# )


Curious, why don't we remove existing argmax_argmin_strategy and enable this strategy registration for those ops? This happens to others PRs in the stack also.

I have a PR at the end for enablement of new strategies and deleting existing strategies. The reason I do it this way is so that if there are issues with logic, there's only one pr reverted, making it cleaner for me and others who are also working on this file

zpcore

The logic LGTM comparing with the existing strategy.

anshul-si · 2026-04-14T17:11:21Z

@pytorchbot review this pr

pytorch-bot · 2026-04-14T17:11:24Z

❌ 🤖 pytorchbot command failed:

@pytorchbot: error: argument command: invalid choice: 'review' (choose from 'merge', 'revert', 'rebase', 'label', 'drci', 'lint', 'fix-lint', 'apply-lint', 'cherry-pick')

usage: @pytorchbot [-h]
                   
                   {merge,revert,rebase,label,drci,lint,fix-lint,apply-lint,cherry-pick}
                   ...

Try @pytorchbot --help for more info.

…with-indices ops" Refactor max.dim, min.dim, median.dim, nanmedian.dim, mode, kthvalue to use the shared _shard_non_reduction_dim helper. Add new single_dim_strategy for argmax/argmin (decorator commented out, old register_op_strategy kept). [ghstack-poisoned]

anshul-si · 2026-06-11T18:21:29Z

@pytorchbot merge

pytorchmergebot · 2026-06-11T18:24:53Z

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

…s ops (pytorch#179200) Refactor max.dim, min.dim, median.dim, nanmedian.dim, mode, kthvalue to use the shared _shard_non_reduction_dim helper. Add new single_dim_strategy for argmax/argmin (decorator commented out, old register_op_strategy kept). Pull Request resolved: pytorch#179200 Approved by: https://github.com/zpcore ghstack dependencies: pytorch#179068

Add single_dim_strategy implementations for: - Category C: Linear reductions (sum, mean, amax, amin, nansum, prod) - Category C: Bool reductions (all, any) - fixes incorrect S(dim)->P rules - Category D: Non-linear reductions (std, var, var_mean, logsumexp) - Category E: Norm reductions (vector_norm, norm.Scalar, powsum) - fixes incorrect S(dim)->NormP rules - Category J: Global median (median.default, nanmedian.default) New registrations are commented out; old register_op_strategy kept active. Pull Request resolved: #179201 Approved by: https://github.com/zpcore, https://github.com/pianpwk ghstack dependencies: #179068, #179200

…ly, and histc ops (#179202) Add single_dim_strategy implementations for: - Category G: Pooling ops (batch/channel dim sharding, partial for avg pools) - Category H: Replicate-only ops (SVD, QR, diagonal, tril, triu, upsample, etc.) - Category L: Histc (partial output when min/max specified) New registrations are commented out; old register_op_strategy kept active. Pull Request resolved: #179202 Approved by: https://github.com/pianpwk ghstack dependencies: #179068, #179200, #179201

…s ops (pytorch#179200) Refactor max.dim, min.dim, median.dim, nanmedian.dim, mode, kthvalue to use the shared _shard_non_reduction_dim helper. Add new single_dim_strategy for argmax/argmin (decorator commented out, old register_op_strategy kept). Pull Request resolved: pytorch#179200 Approved by: https://github.com/zpcore ghstack dependencies: pytorch#179068

…h#179201) Add single_dim_strategy implementations for: - Category C: Linear reductions (sum, mean, amax, amin, nansum, prod) - Category C: Bool reductions (all, any) - fixes incorrect S(dim)->P rules - Category D: Non-linear reductions (std, var, var_mean, logsumexp) - Category E: Norm reductions (vector_norm, norm.Scalar, powsum) - fixes incorrect S(dim)->NormP rules - Category J: Global median (median.default, nanmedian.default) New registrations are commented out; old register_op_strategy kept active. Pull Request resolved: pytorch#179201 Approved by: https://github.com/zpcore, https://github.com/pianpwk ghstack dependencies: pytorch#179068, pytorch#179200

…ly, and histc ops (pytorch#179202) Add single_dim_strategy implementations for: - Category G: Pooling ops (batch/channel dim sharding, partial for avg pools) - Category H: Replicate-only ops (SVD, QR, diagonal, tril, triu, upsample, etc.) - Category L: Histc (partial output when min/max specified) New registrations are commented out; old register_op_strategy kept active. Pull Request resolved: pytorch#179202 Approved by: https://github.com/pianpwk ghstack dependencies: pytorch#179068, pytorch#179200, pytorch#179201

…s ops (pytorch#179200) Refactor max.dim, min.dim, median.dim, nanmedian.dim, mode, kthvalue to use the shared _shard_non_reduction_dim helper. Add new single_dim_strategy for argmax/argmin (decorator commented out, old register_op_strategy kept). Pull Request resolved: pytorch#179200 Approved by: https://github.com/zpcore ghstack dependencies: pytorch#179068

…h#179201) Add single_dim_strategy implementations for: - Category C: Linear reductions (sum, mean, amax, amin, nansum, prod) - Category C: Bool reductions (all, any) - fixes incorrect S(dim)->P rules - Category D: Non-linear reductions (std, var, var_mean, logsumexp) - Category E: Norm reductions (vector_norm, norm.Scalar, powsum) - fixes incorrect S(dim)->NormP rules - Category J: Global median (median.default, nanmedian.default) New registrations are commented out; old register_op_strategy kept active. Pull Request resolved: pytorch#179201 Approved by: https://github.com/zpcore, https://github.com/pianpwk ghstack dependencies: pytorch#179068, pytorch#179200

…ly, and histc ops (pytorch#179202) Add single_dim_strategy implementations for: - Category G: Pooling ops (batch/channel dim sharding, partial for avg pools) - Category H: Replicate-only ops (SVD, QR, diagonal, tril, triu, upsample, etc.) - Category L: Histc (partial output when min/max specified) New registrations are commented out; old register_op_strategy kept active. Pull Request resolved: pytorch#179202 Approved by: https://github.com/pianpwk ghstack dependencies: pytorch#179068, pytorch#179200, pytorch#179201

pytorch-bot Bot added ciflow/dtensor Run DTensor specific tests ciflow/inductor ciflow/torchtitan Run TorchTitan integration tests release notes: distributed (dtensor) release notes category labels Apr 2, 2026

anshul-si requested review from pianpwk, wconstab and zpcore April 3, 2026 02:05

anshul-si added 2 commits April 2, 2026 20:27

zpcore reviewed Apr 14, 2026

View reviewed changes

zpcore approved these changes Apr 14, 2026

View reviewed changes

anshul-si added the ciflow/trunk Trigger trunk jobs on your pull request label Apr 14, 2026

anshul-si added 6 commits April 14, 2026 10:18

anshul-si added 6 commits June 10, 2026 14:42

pytorchmergebot added the merging label Jun 11, 2026

pytorchmergebot closed this in 26e47bd Jun 11, 2026

pytorchmergebot added Merged and removed merging labels Jun 11, 2026

github-actions Bot deleted the gh/anshul-si/125/head branch July 12, 2026 02:25

Navigation Menu

Repository navigation

[dtensor] Add single_dim_strategy functions for reduction-with-indices ops - #179200

[dtensor] Add single_dim_strategy functions for reduction-with-indices ops#179200
anshul-si wants to merge 16 commits into
gh/anshul-si/125/basefrom
gh/anshul-si/125/head

anshul-si commented Apr 2, 2026 •

edited

Loading

Uh oh!

pytorch-bot Bot commented Apr 2, 2026 •

edited

Loading

Uh oh!

zpcore Apr 14, 2026

Uh oh!

anshul-si Apr 14, 2026

Uh oh!

zpcore left a comment

Uh oh!

anshul-si commented Apr 14, 2026

Uh oh!

pytorch-bot Bot commented Apr 14, 2026

Uh oh!

anshul-si commented Jun 11, 2026

Uh oh!

pytorchmergebot commented Jun 11, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

anshul-si commented Apr 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot Bot commented Apr 2, 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/179200

✅ You can merge normally! (5 Unrelated Failures)

Uh oh!

zpcore Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

anshul-si Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

zpcore left a comment

Choose a reason for hiding this comment

Uh oh!

anshul-si commented Apr 14, 2026

Uh oh!

pytorch-bot Bot commented Apr 14, 2026

Uh oh!

anshul-si commented Jun 11, 2026

Uh oh!

pytorchmergebot commented Jun 11, 2026

Merge started

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

anshul-si commented Apr 2, 2026 •

edited

Loading

pytorch-bot Bot commented Apr 2, 2026 •

edited

Loading