8000 Support `torch.linalg.trace` by asi1024 · Pull Request #62714 · pytorch/pytorch · GitHub
[go: up one dir, main page]
Skip to content

Support torch.linalg.trace #62714

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 23 commits into from
Closed

Support torch.linalg.trace #62714

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
23 commits
Select commit Hold shift + click to select a range
0b91b4b
Add torch.linalg.trace
asi1024 Aug 4, 2021
ee1901c
Add tests
asi1024 Sep 13, 2021
d8a71b4
Fix
asi1024 Sep 21, 2021
467580b
Fix linalg.trace docstring
asi1024 Sep 29, 2021
361acf5
refactor tests
asi1024 Sep 29, 2021
2edb1c8
Fix linalg.trace backward failure
asi1024 Sep 29, 2021
b9a6a7e
Remove CompositeExplicitAutograd
asi1024 Sep 29, 2021
4e96d97
Fix doc
asi1024 Sep 29, 2021
5858e22
Fix
asi1024 Sep 29, 2021
c7061ab
Merge branch 'master' into trace
asi1024 Jan 5, 2022
f98c8e9
Support out argument
asi1024 Jan 5, 2022
cecc545
test
asi1024 Jan 11, 2022
5f073d3
Merge remote-tracking branch 'pytorch/master' into trace
asi1024 Jan 11, 2022
37914e4
Merge remote-tracking branch 'pytorch/master' into trace
asi1024 Jan 31, 2022
aa1b23f
Fix
asi1024 Feb 1, 2022
025c6e2
Update trace documentation
asi1024 Feb 3, 2022
de21898
Fix test
asi1024 Feb 10, 2022
9986e30
Merge branch 'master' into trace
asi1024 Feb 13, 2022
cabf8dd
Update linalg.rst
asi1024 Feb 15, 2022
bf17c76
Fix for doctest
asi1024 Feb 19, 2022
d857bd8
Merge remote-tracking branch 'pytorch/master' into trace
asi1024 Feb 20, 2022
e5872eb
Fix tests
asi1024 Feb 24, 2022
9622b99
Fix input name
asi1024 Feb 25, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Fix tests
  • Loading branch information
@asi1024
asi1024 committed Feb 25, 2022
commit e5872eb7515dc39b2aadcb7e2f32bea2919db759
20 changes: 19 additions & 1 deletion test/test_linalg.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7776,7 +7776,7 @@ def test_linalg_trace(self, device):
for args in inputs:
shape = args['shape']
size = np.prod(shape)
x = torch.arange(size, device=device).reshape(*shape)
x = torch.arange(size, device=device).reshape(shape)
y = torch.linalg.trace(x)
yn = np.trace(x.cpu(), axis1=-2, axis2=-1)
self.assertEqual(y, yn)
Expand All @@ -7786,6 +7786,24 @@ def test_linalg_trace(self, device):
yn = np.trace(x.cpu(), offset, axis1=-2, axis2=-1)
self.assertEqual(y, yn)

def test_linalg_trace_error(self, device):
inputs = [
# Less than 2-dim
{'shape': (), 'offsets': [0]},
{'shape': (10), 'offsets': [0]},
# offset out of range
{'shape': (0, 0), 'offsets': [0]},
{'shape': (10, 20, 30), 'offsets': [-20, 30]},
]

for args in inputs:
shape = args['shape']
size = np.prod(shape)
x = torch.arange(size, device=device).reshape(shape)
for offset in args['offsets']:
with self.assertRaises(RuntimeError):
torch.linalg.trace(x, offset=offset)

@onlyCUDA
@skipCUDAIfNoMagma
@skipCUDAIfNoCusolver
Expand Down
6 changes: 5 additions & 1 deletion torch/testing/_internal/common_methods_invocations.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2575,10 +2575,14 @@ def sample_inputs_trace(self, device, dtype, requires_grad, **kwargs):
requires_grad=requires_grad))),)

def sample_inputs_linalg_trace(self, device, dtype, requires_grad, **kwargs):
# input.shape, offset
inputs = (
((S, S), 0),
Copy link
Collaborator
@mruberry mruberry Feb 22, 2022 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What happens on empty tensors? We should probably add a case for them. What about a batched sample input too? (S, S, S)?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The trace function implemented in this PR returns different values from numpy.trace for 3-dim inputs. numpy.trace reduces with axis1=0, axis2=1 whereas array API specifies to reduce with axis1=-2, axis2=-1.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it possible to compare then against a lambda with our same defaults that the calls into np.trace?

((S, M), 0),
((S, S), 1),
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Add a sample with a negative offset and a comment explaining the format of these tuples

((S, S), -1),
((S, S, S), 0),
((S, M, L), -2),
)
samples = []

Expand Down Expand Up @@ -9994,7 +9998,7 @@ def ref_pairwise_distance(input1, input2):
sample_inputs_func=sample_inputs_linalg_slogdet,
decorators=[skipCUDAIfNoMagma, skipCPUIfNoLapack],),
OpInfo('linalg.trace',
ref=np.trace,
ref=lambda x, offset=0: np.trace(x, offset, -2, -1),
aten_name='linalg_trace',
op=torch.linalg.trace,
dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.float16),
Expand Down
0