Fix output tensor shape for argmin and argmax where keepdim=True and dim=None by mrnikwaws · Pull Request #6536 · pytorch/xla

mrnikwaws · 2024-02-15T00:04:54Z

Current failure looks like this:

root@f825750bf417:/ansible/torch_xla/pytorch/xla# export PJRT_DEVICE="CPU"
root@f825750bf417:/ansible/torch_xla/pytorch/xla# python
Python 3.8.18 (default, Feb  1 2024, 06:10:58) 
[GCC 10.2.1 20210110] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> import torch_xla
>>> t = torch.rand((3,4))
>>> t
tensor([[0.9253, 0.5503, 0.4175, 0.7273],
        [0.8306, 0.7907, 0.2054, 0.5639],
        [0.4089, 0.1673, 0.9702, 0.4839]])
>>> t_xla = t.to('xla')
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1707952093.212406   88634 cpu_client.cc:404] TfrtCpuClient created.
>>> t_xla
tensor([[0.9253, 0.5503, 0.4175, 0.7273],
        [0.8306, 0.7907, 0.2054, 0.5639],
        [0.4089, 0.1673, 0.9702, 0.4839]], device='xla:0')
>>> torch.argmax(t,dim=None,keepdim=True)
tensor([[10]])
>>> torch.argmax(t_xla,dim=None,keepdim=True)
tensor(10, device='xla:0'

After:

root@f825750bf417:/ansible/torch_xla/pytorch/xla# python
Python 3.8.18 (default, Feb  1 2024, 06:10:58) 
[GCC 10.2.1 20210110] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> import torch_xla
>>> t = torch.randn((3,4))
>>> t
tensor([[-1.6839e+00, -5.5569e-01, -1.1452e+00,  4.5730e-01],
        [ 7.5517e-01,  2.3971e+00,  5.8805e-01,  8.4879e-01],
        [-3.3246e-04,  1.4524e-01,  2.0454e-01, -5.7229e-01]])
>>> t_xla = t.to('xla')
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
I0000 00:00:1707953861.897586   92883 cpu_client.cc:404] TfrtCpuClient created.
>>> torch.argmax(t,keepdim=True)
tensor([[5]])
>>> torch.argmax(t_xla,keepdim=True)
tensor([[5]], device='xla:0')
>>> torch.argmin(t,keepdim=True)
tensor([[0]])
>>> torch.argmin(t_xla,keepdim=True)
tensor([[0]], device='xla:0')

…op_name

…n and torch.argmax

wonjoo-wj

Thanks! Changes LGTM, can you add a corresponding unit test at https://github.com/pytorch/xla/blob/master/test/cpp/test_aten_xla_tensor_2.cpp#L2077 and fix the linter issues?

mrnikwaws · 2024-02-27T19:13:13Z

Yep will do

wonjoo-wj

Thanks! Changes LGTM, let's wait for the CI to verify the tests.

wonjoo-wj · 2024-02-28T00:24:15Z

Seems like the build is failing:

test/cpp/test_aten_xla_tensor_2.cpp:2189:8: note: ‘virtual void torch_xla::cpp_test::AtenXlaTensorTest_TestArgMaxDimKeep_Test::TestBody()’ previously defined here
 2189 | TEST_F(AtenXlaTensorTest, TestArgMaxDimKeep) {

My guess is because test with name TestArgMaxDimKeep already exists.

mrnikwaws · 2024-02-28T00:37:43Z

I misread the test code - they need unique names - fixing now

mrnikwaws · 2024-02-28T17:34:43Z

New test failures seem unrelated?

[0 / 1] [Prepa] BazelWorkspaceStatusAction stable-status.txt
[11 / 13] 7 / 8 tests; Compiling test/cpp/test_aten_xla_tensor_2.cpp; 0s local
[11 / 13] 7 / 8 tests; Compiling test/cpp/test_aten_xla_tensor_2.cpp; 11s local
[12 / 13] 7 / 8 tests; checking cached actions
[12 / 13] 7 / 8 tests; Linking test/cpp/test_aten_xla_tensor_2; 4s local
[12 / 13] 7 / 8 tests; Linking test/cpp/test_aten_xla_tensor_2; 10s local
[13 / 14] 7 / 8 tests; [Prepa] Testing //test/cpp:test_aten_xla_tensor_2
[13 / 14] 7 / 8 tests; Testing //test/cpp:test_aten_xla_tensor_2; 11s local
[13 / 14] 8 / 8 tests; Testing //test/cpp:test_aten_xla_tensor_2; 24s local
INFO: Elapsed time: 85.857s, Critical Path: 83.49s
INFO: 7 processes: 4 internal, 3 local.
INFO: Build completed successfully, 7 total actions
//torch_xla/csrc/runtime:cache_test                             (cached) PASSED in 0.7s
//torch_xla/csrc/runtime:env_hash_test                          (cached) PASSED in 0.6s
//torch_xla/csrc/runtime:ifrt_computation_client_test           (cached) PASSED in 1.2s
//torch_xla/csrc/runtime:pjrt_computation_client_test           (cached) PASSED in 0.5s
//torch_xla/csrc/runtime:sys_util_test                          (cached) PASSED in 0.0s
//torch_xla/csrc/runtime:util_test                              (cached) PASSED in 0.1s
//torch_xla/csrc/runtime:xla_util_test                          (cached) PASSED in 0.8s
//test/cpp:test_aten_xla_tensor_2                                        PASSED in 24.9s

Executed 1 out of 8 tests: 8 tests pass.

wonjoo-wj · 2024-02-29T00:48:50Z

Apologies for the delay but can you rebase with this head? The failing tests have been disabled on head, but I hope to let the CI complete before merging this. Thanks!

mrnikwaws · 2024-02-29T21:27:42Z

Looks like new unrelated failures

wonjoo-wj · 2024-02-29T21:36:27Z

Yeah, it's been a rough day for the head CI.. 😢

Let's just wait for the other CIs and merge it if they're green, changes LGTM anyways. Thanks!

wonjoo-wj · 2024-02-29T22:31:32Z

CI seems to be passing other than the same failure on head. Merging this. Thanks!

…dim=None (pytorch#6536)

mrnikwaws and others added 4 commits December 8, 2023 00:09

Remove redundant log message when using IR debug, but not overriding …

57605e1

…op_name

Merge branch 'pytorch:master' into master

a65a879

Fix the valid use case where dim=None and keepdim=True on torch.argmi…

4184bd9

…n and torch.argmax

Merge branch 'pytorch:master' into master

6735ac1

JackCaoG requested a review from wonjoo-wj February 15, 2024 19:21

wonjoo-wj requested changes Feb 20, 2024

View reviewed changes

mrnikwaws added 2 commits February 27, 2024 19:15

Linter fix

a8bbdb7

Add C++ tests for argmin and argmax with dim=None and keepdim=true

97bd4df

mrnikwaws requested a review from wonjoo-wj February 27, 2024 19:30

wonjoo-wj approved these changes Feb 27, 2024

View reviewed changes

Give tests unique names

1ac2622

Merge branch 'pytorch:master' into master

7a88949

wonjoo-wj merged commit a1ab7fd into pytorch:master Feb 29, 2024

amithrm pushed a commit to amithrm/xla that referenced this pull request Mar 1, 2024

Fix output tensor shape for argmin and argmax where keepdim=True and …

17fd138

…dim=None (pytorch#6536)

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Fix output tensor shape for argmin and argmax where keepdim=True and dim=None#6536

Fix output tensor shape for argmin and argmax where keepdim=True and dim=None#6536
wonjoo-wj merged 8 commits intopytorch:masterfrom
mrnikwaws:master

mrnikwaws commented Feb 15, 2024

Uh oh!

wonjoo-wj left a comment

Uh oh!

mrnikwaws commented Feb 27, 2024

Uh oh!

wonjoo-wj left a comment •

edited

Loading

Uh oh!

wonjoo-wj commented Feb 28, 2024

Uh oh!

mrnikwaws commented Feb 28, 2024

Uh oh!

mrnikwaws commented Feb 28, 2024

Uh oh!

wonjoo-wj commented Feb 29, 2024

Uh oh!

mrnikwaws commented Feb 29, 2024

Uh oh!

wonjoo-wj commented Feb 29, 2024

Uh oh!

wonjoo-wj commented Feb 29, 2024

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

mrnikwaws commented Feb 15, 2024

Uh oh!

wonjoo-wj left a comment

Choose a reason for hiding this comment

Uh oh!

mrnikwaws commented Feb 27, 2024

Uh oh!

wonjoo-wj 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!

wonjoo-wj commented Feb 28, 2024

Uh oh!

mrnikwaws commented Feb 28, 2024

Uh oh!

mrnikwaws commented Feb 28, 2024

Uh oh!

wonjoo-wj commented Feb 29, 2024

Uh oh!

mrnikwaws commented Feb 29, 2024

Uh oh!

wonjoo-wj commented Feb 29, 2024

Uh oh!

wonjoo-wj commented Feb 29, 2024

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

wonjoo-wj left a comment •

edited

Loading