Add `torch.put_along_dim` and `torch.put_along_dim_` like `np.put_along_axis` #125601

arunppsg · 2024-05-06T17:09:22Z

Fixes #120209

cc @albanD

pytorch-bot · 2024-05-06T17:09:25Z

🔗 Helpful Links

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

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

❌ 51 New Failures

As of commit bbcf254 with merge base 60bbdc0 ():

NEW FAILURES - The following jobs have failed:

pull / linux-focal-cuda12.1-py3.10-gcc9 / test (default, 1, 5, linux.4xlarge.nvidia.gpu) (gh)
test_ops_fwd_gradients.py::TestFwdGradientsCUDA::test_fn_fwgrad_bwgrad_put_along_dim_cuda_complex128
pull / linux-focal-cuda12.1-py3.10-gcc9 / test (default, 4, 5, linux.4xlarge.nvidia.gpu) (gh)
functorch/test_vmap.py::TestVmapOperatorsOpInfoCUDA::test_op_has_batch_rule_put_along_dim_cuda_float32
pull / linux-focal-cuda12.1-py3.10-gcc9 / test (default, 5, 5, linux.4xlarge.nvidia.gpu) (gh)
test_ops_gradients.py::TestBwdGradientsCUDA::test_fn_grad_put_along_dim_cuda_complex128
pull / linux-focal-cuda12.1-py3.10-gcc9-experimental-split-build / test (default, 1, 5, linux.4xlarge.nvidia.gpu) (gh)
Process completed with exit code 9.
pull / linux-focal-cuda12.1-py3.10-gcc9-experimental-split-build / test (default, 2, 5, linux.4xlarge.nvidia.gpu) (gh)
Process completed with exit code 9.
pull / linux-focal-cuda12.1-py3.10-gcc9-experimental-split-build / test (default, 3, 5, linux.4xlarge.nvidia.gpu) (gh)
Process completed with exit code 9.
pull / linux-focal-cuda12.1-py3.10-gcc9-experimental-split-build / test (default, 4, 5, linux.4xlarge.nvidia.gpu) (gh)
Process completed with exit code 9.
pull / linux-focal-cuda12.1-py3.10-gcc9-experimental-split-build / test (default, 5, 5, linux.4xlarge.nvidia.gpu) (gh)
Process completed with exit code 9.
pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 1, 5, linux.g5.4xlarge.nvidia.gpu) (gh)
functorch/test_vmap.py::TestVmapOperatorsOpInfoCUDA::test_op_has_batch_rule_put_along_dim_cuda_float32
pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 2, 5, linux.g5.4xlarge.nvidia.gpu) (gh)
test_ops_fwd_gradients.py::TestFwdGradientsCUDA::test_fn_fwgrad_bwgrad_put_along_dim_cuda_complex128
pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 3, 5, linux.g5.4xlarge.nvidia.gpu) (gh)
inductor/test_torchinductor_opinfo.py::TestInductorOpInfoCUDA::test_comprehensive_put_along_dim_cuda_float16
pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 4, 5, linux.g5.4xlarge.nvidia.gpu) (gh)
inductor/test_torchinductor_opinfo.py::TestInductorOpInfoCUDA::test_comprehensive_put_along_dim_cuda_float64
pull / linux-focal-cuda12.1-py3.10-gcc9-sm86 / test (default, 5, 5, linux.g5.4xlarge.nvidia.gpu) (gh)
inductor/test_torchinductor_opinfo.py::TestInductorOpInfoCUDA::test_comprehensive_put_along_dim_cuda_float32
pull / linux-focal-py3.11-clang10 / test (crossref, 1, 2, linux.2xlarge) (gh)
functorch/test_vmap.py::TestVmapOperatorsOpInfoCPU::test_op_has_batch_rule_put_along_dim_cpu_float32
pull / linux-focal-py3.11-clang10 / test (crossref, 2, 2, linux.2xlarge) (gh)
test_proxy_tensor.py::TestProxyTensorOpInfoCPU::test_make_fx_symbolic_exhaustive_out_put_along_dim_cpu_float32
pull / linux-focal-py3.11-clang10 / test (default, 1, 4, linux.2xlarge) (gh)
inductor/test_torchinductor_opinfo.py::TestInductorOpInfoCPU::test_comprehensive_put_along_dim_cpu_float32
pull / linux-focal-py3.11-clang10 / test (default, 2, 4, linux.2xlarge) (gh)
inductor/test_torchinductor_opinfo.py::TestInductorOpInfoCPU::test_comprehensive_put_along_dim_cpu_bool
pull / linux-focal-py3.11-clang10 / test (default, 3, 4, linux.2xlarge) (gh)
inductor/test_torchinductor_opinfo.py::TestInductorOpInfoCPU::test_comprehensive_put_along_dim_cpu_int64
pull / linux-focal-py3.11-clang10 / test (default, 4, 4, linux.2xlarge) (gh)
inductor/test_torchinductor_opinfo.py::TestInductorOpInfoCPU::test_comprehensive_put_along_dim_cpu_int32
pull / linux-focal-py3.11-clang10 / test (dynamo, 1, 3, linux.2xlarge) (gh)
RuntimeError: test_overrides 1/1 failed!
pull / linux-focal-py3.11-clang10 / test (dynamo, 2, 3, linux.2xlarge) (gh)
functorch/test_vmap_registrations.py::TestFunctorchDispatcher::test_unimplemented_batched_registrations_[aten::put_along_dim]
pull / linux-focal-py3.11-clang10 / test (dynamo, 3, 3, linux.2xlarge) (gh)
test_indexing.py::TestIndexingCPU::test_put_along_dim_cpu_float32
pull / linux-focal-py3.12-clang10 / test (default, 1, 4, linux.2xlarge) (gh)
inductor/test_torchinductor_opinfo.py::TestInductorOpInfoCPU::test_comprehensive_put_along_dim_cpu_float64
pull / linux-focal-py3.12-clang10 / test (default, 2, 4, linux.2xlarge) (gh)
inductor/test_torchinductor_opinfo.py::TestInductorOpInfoCPU::test_comprehensive_put_along_dim_cpu_int32
pull / linux-focal-py3.12-clang10 / test (default, 3, 4, linux.2xlarge) (gh)
inductor/test_torchinductor_opinfo.py::TestInductorOpInfoCPU::test_comprehensive_put_along_dim_cpu_float16
pull / linux-focal-py3.12-clang10 / test (dynamo, 1, 3, linux.2xlarge) (gh)
RuntimeError: test_overrides 1/1 failed!
pull / linux-focal-py3.12-clang10 / test (dynamo, 3, 3, linux.2xlarge) (gh)
test_indexing.py::TestIndexingCPU::test_put_along_dim_cpu_float32
pull / linux-focal-py3.12-clang10-experimental-split-build / test (default, 1, 3, linux.2xlarge) (gh)
Process completed with exit code 9.
pull / linux-focal-py3.12-clang10-experimental-split-build / test (default, 2, 3, linux.2xlarge) (gh)
Process completed with exit code 9.
pull / linux-focal-py3.12-clang10-experimental-split-build / test (default, 3, 3, linux.2xlarge) (gh)
Process completed with exit code 9.
pull / linux-focal-py3.12-clang10-experimental-split-build / test (dynamo, 1, 3, linux.2xlarge) (gh)
Process completed with exit code 9.
pull / linux-focal-py3.12-clang10-experimental-split-build / test (dynamo, 2, 3, linux.2xlarge) (gh)
Process completed with exit code 9.
pull / linux-focal-py3.12-clang10-experimental-split-build / test (dynamo, 3, 3, linux.2xlarge) (gh)
Process completed with exit code 9.
pull / linux-focal-py3.8-clang10 / test (crossref, 1, 2, linux.2xlarge) (gh)
functorch/test_vmap.py::TestVmapOperatorsOpInfoCPU::test_op_has_batch_rule_put_along_dim_cpu_float32
pull / linux-focal-py3.8-clang10 / test (crossref, 2, 2, linux.2xlarge) (gh)
test_proxy_tensor.py::TestProxyTensorOpInfoCPU::test_make_fx_symbolic_exhaustive_out_put_along_dim_cpu_float32
pull / linux-focal-py3.8-clang10 / test (default, 1, 4, linux.2xlarge) (gh)
inductor/test_torchinductor_opinfo.py::TestInductorOpInfoCPU::test_comprehensive_put_along_dim_cpu_float64
pull / linux-focal-py3.8-clang10 / test (default, 2, 4, linux.2xlarge) (gh)
inductor/test_torchinductor_opinfo.py::TestInductorOpInfoCPU::test_comprehensive_put_along_dim_cpu_int32
pull / linux-focal-py3.8-clang10 / test (default, 3, 4, linux.2xlarge) (gh)
inductor/test_torchinductor_opinfo.py::TestInductorOpInfoCPU::test_comprehensive_put_along_dim_cpu_float16
pull / linux-focal-py3.8-clang10 / test (default, 4, 4, linux.2xlarge) (gh)
test_ops_fwd_gradients.py::TestFwdGradientsCPU::test_fn_fwgrad_bwgrad_put_along_dim_cpu_complex128
pull / linux-focal-py3.8-clang10 / test (dynamo, 1, 3, linux.2xlarge) (gh)
RuntimeError: test_overrides 1/1 failed!
pull / linux-focal-py3.8-clang10 / test (dynamo, 2, 3, linux.2xlarge) (gh)
functorch/test_vmap_registrations.py::TestFunctorchDispatcher::test_unimplemented_batched_registrations_[aten::put_along_dim]
pull / linux-focal-py3.8-clang10 / test (dynamo, 3, 3, linux.2xlarge) (gh)
test_indexing.py::TestIndexingCPU::test_put_along_dim_cpu_float32
pull / linux-focal-py3.8-clang10-onnx / test (default, 1, 2, linux.2xlarge) (gh)
onnx/test_fx_op_consistency.py::TestOnnxModelOutputConsistency_opset_version_18_model_type_TorchModelType.TORCH_NN_MODULECPU::test_output_match_put_along_dim_cpu_bool
pull / linux-focal-py3.8-clang10-onnx / test (default, 2, 2, linux.2xlarge) (gh)
onnx/test_fx_op_consistency.py::TestOnnxModelOutputConsistency_opset_version_18_model_type_TorchModelType.TORCH_NN_MODULECPU::test_output_match_put_along_dim_cpu_int32
pull / linux-jammy-py3.10-clang15-asan / test (default, 4, 6, linux.4xlarge) (gh)
test_ops_fwd_gradients.py::TestFwdGradientsCPU::test_fn_fwgrad_bwgrad_put_along_dim_cpu_complex128
pull / linux-jammy-py3.10-clang15-asan / test (default, 5, 6, linux.4xlarge) (gh)
test_proxy_tensor.py::TestProxyTensorOpInfoCPU::test_make_fx_symbolic_exhaustive_out_put_along_dim_cpu_float32
pull / linux-jammy-py3.10-clang15-asan / test (default, 6, 6, linux.4xlarge) (gh)
functorch/test_vmap.py::TestVmapOperatorsOpInfoCPU::test_op_has_batch_rule_put_along_dim_cpu_float32
pull / linux-jammy-py3.8-gcc11 / test (default, 1, 4, linux.2xlarge) (gh)
inductor/test_torchinductor_opinfo.py::TestInductorOpInfoCPU::test_comprehensive_put_along_dim_cpu_int32
pull / linux-jammy-py3.8-gcc11 / test (default, 2, 4, linux.2xlarge) (gh)
inductor/test_torchinductor_opinfo.py::TestInductorOpInfoCPU::test_comprehensive_put_along_dim_cpu_float16
pull / linux-jammy-py3.8-gcc11 / test (default, 3, 4, linux.2xlarge) (gh)
inductor/test_torchinductor_opinfo.py::TestInductorOpInfoCPU::test_comprehensive_put_along_dim_cpu_bool
pull / linux-jammy-py3.8-gcc11 / test (default, 4, 4, linux.2xlarge) (gh)
inductor/test_torchinductor_opinfo.py::TestInductorOpInfoCPU::test_comprehensive_put_along_dim_cpu_float32

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

lezcano

This looks quite good!

Now, rather than adding a one-off test, could you add an OpInfo within common_methods_invocations.py. You can even specify a NumPy reference there and a test will test your implementation against it.

Make sure all the relevant tests in test_ops*.py pass locally before submitting :)

aten/src/ATen/native/TensorAdvancedIndexing.cpp

lezcano · 2024-06-17T15:26:17Z

test/test_indexing.py

+ np.put_along_axis(t_np, indices_np, values_np, dim)
+ self.assertEqual(actual, t_np, atol=0, rtol=0)
+
+ t = torch.tensor([[10, 30, 20], [60, 40, 50]], dtype=dtype)


can you test the broadcasting behaviour between the different tensors, for both dim specified and =None?

aten/src/ATen/native/TensorAdvancedIndexing.cpp

arunppsg · 2024-06-20T12:12:48Z

Made some of the change and work in progress for OpInfo tests.

lezcano · 2024-06-20T12:35:19Z

changes so far look good. Thank you!

arunppsg · 2024-07-04T17:06:01Z

Added simple OpInfo test, though one test is failing: test_non_standard_bool_values_put_along_dim_cpu_bool - working on debugging it.

pytorchbot added the open source label May 6, 2024

arunppsg added 8 commits June 15, 2024 19:01

added a simple testcase

de4c84d

add put_along_dim to native_functions.yaml

e7b9b8a

added put_along_dim for the simplest case [skip ci]

15ecec8

add more tests

7fde3fc

add minimal support for put_along_dim

8b2899d

added helper function

c698dcc

added support for axis=None case

a8e53ab

add docstring

1f08e48

arunppsg force-pushed the issue-120209 branch from 1a4aca2 to 1f08e48 Compare June 15, 2024 16:55

arunppsg marked this pull request as ready for review June 15, 2024 16:57

mikaylagawarecki requested a review from lezcano June 17, 2024 15:19

mikaylagawarecki added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Jun 17, 2024

lezcano added the module: python frontend For issues relating to PyTorch's Python frontend label Jun 17, 2024

lezcano reviewed Jun 17, 2024

View reviewed changes

arunppsg added 3 commits June 17, 2024 21:39

removed dead code

dcd3400

better new index computation

b5a43f5

replace flatten() with Tensor.put

3f06685

arunppsg added 3 commits June 21, 2024 11:54

place OpInfo(put) in alphabetical order after OpInfo(pow)

5ad0d5a

minor fix

4bdc9ba

add basic op-info tests

bbcf254

arunppsg requested a review from mruberry as a code owner July 4, 2024 16:59

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add `torch.put_along_dim` and `torch.put_along_dim_` like `np.put_along_axis` #125601

Add `torch.put_along_dim` and `torch.put_along_dim_` like `np.put_along_axis` #125601

arunppsg commented May 6, 2024 •

edited by pytorch-bot bot

Loading

pytorch-bot bot commented May 6, 2024 •

edited

Loading

lezcano left a comment

lezcano Jun 17, 2024

arunppsg commented Jun 20, 2024

lezcano commented Jun 20, 2024

arunppsg commented Jul 4, 2024

Add torch.put_along_dim and torch.put_along_dim_ like np.put_along_axis #125601

Are you sure you want to change the base?

Add torch.put_along_dim and torch.put_along_dim_ like np.put_along_axis #125601

Conversation

arunppsg commented May 6, 2024 • edited by pytorch-bot bot Loading

pytorch-bot bot commented May 6, 2024 • edited Loading

🔗 Helpful Links

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

❌ 51 New Failures

lezcano left a comment

Choose a reason for hiding this comment

lezcano Jun 17, 2024

Choose a reason for hiding this comment

arunppsg commented Jun 20, 2024

lezcano commented Jun 20, 2024

arunppsg commented Jul 4, 2024

Add `torch.put_along_dim` and `torch.put_along_dim_` like `np.put_along_axis` #125601

Add `torch.put_along_dim` and `torch.put_along_dim_` like `np.put_along_axis` #125601

arunppsg commented May 6, 2024 •

edited by pytorch-bot bot

Loading

pytorch-bot bot commented May 6, 2024 •

edited

Loading