Support `torch.linalg.trace` by asi1024 · Pull Request #62714 · pytorch/pytorch

asi1024 · 2021-08-04T11:46:10Z

Fixes #62255 (cc/ @mruberry, @rgommers, @emcastillo, @kmaehashi)

This PR adds support of torch.linalg.trace for the compatibility with NumPy's interface and Python array API standard.

>>> torch.linalg.trace(torch.arange(18).reshape(2, 3, 3))
tensor([12, 14, 16])

TODO:

Add documentation
Add tests

cc @jianyuh @nikitaved @pearu @mruberry @walterddr @IvanYashchuk @xwang233 @lezcano

facebook-github-bot · 2021-08-04T11:46:16Z

🔗 Helpful links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/62714
📄 Preview docs built from this PR
📄 Preview C++ docs built from this PR
🔧 Opt-in to CIFlow to control what jobs run on your PRs

💊 CI failures summary and remediations

As of commit 9622b99 (more details on the Dr. CI page):

2/2 failures introduced in this PR

🕵️ 1 new failure recognized by patterns

The following CI failures do not appear to be due to upstream breakages:

win-vs2019-cuda11.3-py3 / test (default, 2, 2, windows.8xlarge.nvidia.gpu) (1/1)

Step: "Test" (full log | diagnosis details | 🔁 rerun)

2022-02-25T11:20:02.4745605Z FAIL [0.019s]: test_sparse_addmm_cpu_bfloat16 (__main__.TestSparseCPU)

2022-02-25T11:20:02.4603374Z   test_sparse_zeros_tanh_cuda_float64 (__main__.TestSparseUnaryUfuncsCUDA) ... ok (0.000s)
2022-02-25T11:20:02.4623876Z   test_sparse_zeros_tanh_cuda_int16 (__main__.TestSparseUnaryUfuncsCUDA) ... ok (0.000s)
2022-02-25T11:20:02.4644090Z   test_sparse_zeros_tanh_cuda_int32 (__main__.TestSparseUnaryUfuncsCUDA) ... ok (0.000s)
2022-02-25T11:20:02.4664636Z   test_sparse_zeros_tanh_cuda_int64 (__main__.TestSparseUnaryUfuncsCUDA) ... ok (0.000s)
2022-02-25T11:20:02.4684392Z   test_sparse_zeros_tanh_cuda_int8 (__main__.TestSparseUnaryUfuncsCUDA) ... ok (0.000s)
2022-02-25T11:20:02.4704217Z   test_sparse_zeros_tanh_cuda_uint8 (__main__.TestSparseUnaryUfuncsCUDA) ... ok (0.000s)
2022-02-25T11:20:02.4724154Z   test_sparse_zeros_trunc_cuda_float32 (__main__.TestSparseUnaryUfuncsCUDA) ... ok (0.016s)
2022-02-25T11:20:02.4743513Z   test_sparse_zeros_trunc_cuda_float64 (__main__.TestSparseUnaryUfuncsCUDA) ... ok (0.000s)
2022-02-25T11:20:02.4744595Z 
2022-02-25T11:20:02.4744990Z ======================================================================
2022-02-25T11:20:02.4745605Z FAIL [0.019s]: test_sparse_addmm_cpu_bfloat16 (__main__.TestSparseCPU)
2022-02-25T11:20:02.4746326Z ----------------------------------------------------------------------
2022-02-25T11:20:02.4747129Z Traceback (most recent call last):
2022-02-25T11:20:02.4749094Z   File "C:\actions-runner\_work\pytorch\pytorch\build\win_tmp\build\torch\testing\_internal\common_device_type.py", line 376, in instantiated_test
2022-02-25T11:20:02.4750339Z     result = test(self, **param_kwargs)
2022-02-25T11:20:02.4751375Z   File "C:\actions-runner\_work\pytorch\pytorch\build\win_tmp\build\torch\testing\_internal\common_utils.py", line 2951, in wrapped
2022-02-25T11:20:02.4752299Z     f(self, *args, **kwargs, coalesced=False)
2022-02-25T11:20:02.4753010Z   File "test_sparse.py", line 1275, in test_sparse_addmm
2022-02-25T11:20:02.4753575Z     test_shape(7, 8, 9, 20, False, (1, 1))
2022-02-25T11:20:02.4754352Z   File "test_sparse.py", line 1264, in test_shape
2022-02-25T11:20:02.4754943Z     self.assertEqual(Y, Y_dense)

1 failure not recognized by patterns:

Job	Step	Action
^{Lint / clang-format}	^{Run clang-format}	🔁 rerun

This comment was automatically generated by Dr. CI (expand for details).

Please report bugs/suggestions to the (internal) Dr. CI Users group.

Click here to manually regenerate this comment.

mruberry · 2021-08-17T07:10:16Z

Hey @asi1024, just checking in on this PR because it's still marked as "draft." Is it ready for a review?

asi1024 · 2021-08-30T06:24:10Z

@mruberry Sorry for my late response. I will mark "ready for review" after adding tests and documentation!

IvanYashchuk

Hey, @asi1024, thanks for your contribution! I left a suggestion to use the CompositeImplicitAutograd dispatch key that would allow us to remove the _backward function trimming down unnecessary code. After that, I think the PR should be good to go.

IvanYashchuk · 2021-09-16T13:09:09Z

This line is overwritten by CompositeExplicitAutograd: linalg_trace. The code in the linalg_trace function is independent of the device so CPU, CUDA specialization is not needed here and CompositeExplicitAutograd is the correct choice of the dispatch key.

Suggested change

CPU, CUDA: linalg_trace

On a second thought using CompositeImplicitAutograd should be better, then the backward function is not needed.

IvanYashchuk · 2021-09-16T13:11:59Z

Could you please remove this entry from native_functions.yaml?
Most of the backward functions in PyTorch are placed in torch\csrc\autograd\FunctionsManual.cpp and torch\csrc\autograd\FunctionsManual.h, so let's move linalg_trace_backward from ReduceOps.cpp.

IvanYashchuk · 2021-09-16T13:15:55Z

Is this change needed in this PR?
The function prod(Tensor self, int dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor should be autogenerated with

pytorch/aten/src/ATen/native/native_functions.yaml

Line 4200 in 54d060a

structured_delegate: prod.int_out

codecov · 2021-09-16T16:14:47Z

Codecov Report

Merging #62714 (32e5ee4) into master (0dc9872) will increase coverage by 0.08%.
The diff coverage is 100.00%.

❗ Current head 32e5ee4 differs from pull request most recent head 7092683. Consider uploading reports for the commit 7092683 to get more accurate results

@@            Coverage Diff             @@
##           master   #62714      +/-   ##
==========================================
+ Coverage   66.37%   66.46%   +0.08%     
==========================================
  Files         738      727      -11     
  Lines       94170    93581     -589     
==========================================
- Hits        62510    62200     -310     
+ Misses      31660    31381     -279

lezcano

Left a few points regarding docs / testing.

lezcano · 2021-09-21T08:15:57Z

What diagonal? Given that we have the parameter offset, this should probably read:

Returns the sum of the elements of a diagonal.

Followed by an explanation of how the offset parameter chooses a diagonal.

lezcano · 2021-09-21T08:20:12Z

The call to torch.from_numpy here and a few lines below is not necessary as assertEqual is able to compare tensors and numpy arrays. Even more, not calling it is often faster. Same below.

lezcano · 2021-09-21T08:24:04Z

This might work without all the explicit castings by simply doing

Suggested change

xn = np.array(x.cpu().numpy()).reshape(shape)

yn = np.trace(xn, axis1=-2, axis2=-1)

yn = np.trace(x.cpu(), axis1=-2, axis2=-1)

Same below.

asi1024 · 2021-09-29T08:21:12Z

@lezcano Thank you for your reviews! Could you take another look?

asi1024 · 2022-02-14T03:11:01Z

@lezcano Now all CIs have passed! PTAL!

lezcano · 2022-02-14T14:09:35Z

Yeah, this looks good to me. I just found (yet another, ugh, sorry).
We're missing to add the relevant entry to the docs in docs/source/linalg.rst! It could probably go between diagonal and det.

Sorry for that! I believe this is the last missing thing! :D
Otherwise, we can just wait for @mruberry to review.

asi1024 · 2022-02-21T04:16:21Z

@lezcano The CI failures look unrelated to this PR. Could you take another look?

lezcano · 2022-02-21T12:36:09Z

As mentioned, this LGTM. We now just need to wait for @mruberry to have a look. He's been a bit busy lately, but let's hope he finds some time soon :)

mruberry · 2022-02-22T19:10:28Z

+trace = _add_docstr(_linalg.linalg_trace, r"""
+trace(input, *, offset=0, out=None) -> Tensor
+
+Computes the trace of a matrix.


This is really well written.

mruberry · 2022-02-22T19:11:40Z

+    inputs = (
+        ((S, S), 0),
+        ((S, M), 0),
+        ((S, S), 1),


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

mruberry · 2022-02-22T19:16:21Z

           sample_inputs_func=sample_inputs_linalg_slogdet,
           decorators=[skipCUDAIfNoMagma, skipCPUIfNoLapack],),
+    OpInfo('linalg.trace',
+           ref=np.trace,


Nice reference

mruberry · 2022-02-22T19:16:57Z

        an = torch.from_numpy(np.tensordot(np.zeros((), dtype=np.float32), np.zeros((), dtype=np.float32), 0))
        self.assertEqual(a, an)

+    def test_linalg_trace(self, device):


mruberry · 2022-02-22T19:17:30Z


+def sample_inputs_linalg_trace(self, device, dtype, requires_grad, **kwargs):
+    inputs = (
+        ((S, S), 0),


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

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.

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

mruberry · 2022-02-22T19:18:19Z

+    def test_linalg_trace(self, device):
+        inputs = [
+            {'shape': (1, 1), 'offsets': [0]},
+            {'shape': (10, 1), 'offsets': [0, -9]},


What happens if offset is an absurd number, like 100?

RuntimeError will be raised if offset is out of range. I will add a test for this case!

RuntimeError will be raised if offset is out of range. I will add a test for this case!

Make it an ErrorInput

mruberry · 2022-02-22T19:18:31Z


+    def test_linalg_trace(self, device):
+        inputs = [
+            {'shape': (1, 1), 'offsets': [0]},


Adding the empty case here would be interesting, too

mruberry · 2022-02-22T19:18:49Z

  device_check: NoCheck
  device_guard: False

+- func: linalg_trace.out(Tensor self, *, int offset=0, Tensor(a!) out) -> Tensor(a!)


Nice schemas

mruberry · 2022-02-22T19:19:51Z

+// see https://github.com/pytorch/pytorch/pull/47305,
+Tensor linalg_trace(const Tensor& self, int64_t offset) {
+  TORCH_CHECK(self.dim() >= 2,
+           "self should have at least 2 dimensions, but has ", self.dim(), " dimensions instead");


The user documented name is input (per your docs below) and these warnings should start with the name of the operation like this:

torch.linalg.trace(): input should have at least...

It might be nice to change the user-facing name of this argument to A, which is the name we use throughout torch.linalg

mruberry

Hey @asi1024! Overall this looks really good but there are a few comments inline from me and @lezcano that still need to be addressed. Just ping me when they're done and we'll get this merged.

asi1024 · 2022-02-26T08:46:53Z

@mruberry Updated tests. PTAL!

github-actions · 2022-06-12T10:35:29Z

Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as Stale.
Feel free to remove the Stale label if you feel this was a mistake.
If you are unable to remove the Stale label please contact a maintainer in order to do so.
If you want the bot to never mark this PR stale again, add the no-stale label.
Stale pull requests will automatically be closed after 30 days of inactivity.

facebook-github-bot added the cla signed label Aug 4, 2021

pytorchbot added the open source label Aug 4, 2021

nateanl mentioned this pull request Aug 17, 2021

Add Basic MVDR module pytorch/audio#1708

Merged

5 tasks

asi1024 force-pushed the trace branch from 2c42d7c to 32e5ee4 Compare September 16, 2021 09:51

asi1024 marked this pull request as ready for review September 16, 2021 09:52

asi1024 requested review from IvanYashchuk, albanD, ezyang, lezcano, nikitaved and soulitzer as code owners September 16, 2021 09:52

asi1024 changed the title ~~[WIP] Support torch.linalg.trace~~ Support torch.linalg.trace Sep 16, 2021

IvanYashchuk reviewed Sep 16, 2021

View reviewed changes

ezyang removed their request for review September 16, 2021 14:35

soulitzer removed their request for review September 16, 2021 16:53

albanD removed their request for review September 16, 2021 17:47

heitorschueroff added module: linear algebra Issues related to specialized linear algebra operations in PyTorch; includes matrix multiply matmul triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module labels Sep 16, 2021

asi1024 force-pushed the trace branch from b259ec4 to 7092683 Compare September 21, 2021 07:07

lezcano reviewed Sep 21, 2021

View reviewed changes

asi1024 added 5 commits September 29, 2021 11:55

Add torch.linalg.trace

0b91b4b

Add tests

ee1901c

Fix

d8a71b4

Fix linalg.trace docstring

467580b

refactor tests

361acf5

Update trace documentation

025c6e2

nateanl mentioned this pull request Feb 9, 2022

Add psd and mvdr methods to functional & Refactor PSD and MVDR module in transforms pytorch/audio#2181

Closed

asi1024 added 2 commits February 10, 2022 06:44

Fix test

de21898

Merge branch 'master' into trace

9986e30

asi1024 added 3 commits February 15, 2022 05:55

Update linalg.rst

cabf8dd

Fix for doctest

bf17c76

Merge remote-tracking branch 'pytorch/master' into trace

d857bd8

mruberry reviewed Feb 22, 2022

View reviewed changes

mruberry self-requested a review February 22, 2022 19:21

mruberry reviewed Feb 22, 2022

View reviewed changes

asi1024 added 2 commits February 25, 2022 05:38

Fix tests

e5872eb

Fix input name

9622b99

suo removed the ciflow/default label Mar 22, 2022

rgommers added the module: python array api Issues related to the Python Array API label Apr 13, 2022

github-actions Bot added the Stale label Jun 12, 2022

github-actions Bot closed this Jul 12, 2022

	xn = np.array(x.cpu().numpy()).reshape(shape)
	yn = np.trace(xn, axis1=-2, axis2=-1)
	yn = np.trace(x.cpu(), axis1=-2, axis2=-1)

Conversation

asi1024 commented Aug 4, 2021 • edited by pytorch-probot Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

facebook-github-bot commented Aug 4, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful links

💊 CI failures summary and remediations

🕵️ 1 new failure recognized by patterns

win-vs2019-cuda11.3-py3 / test (default, 2, 2, windows.8xlarge.nvidia.gpu) (1/1)

1 failure not recognized by patterns:

Uh oh!

mruberry commented Aug 17, 2021

Uh oh!

asi1024 commented Aug 30, 2021

Uh oh!

IvanYashchuk left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

codecov Bot commented Sep 16, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

lezcano left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

asi1024 commented Sep 29, 2021

Uh oh!

asi1024 commented Feb 14, 2022

Uh oh!

lezcano commented Feb 14, 2022

Uh oh!

asi1024 commented Feb 21, 2022

Uh oh!

lezcano commented Feb 21, 2022

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mruberry Feb 22, 2022 • 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!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

asi1024 commented Aug 4, 2021 •

edited by pytorch-probot Bot

Loading

facebook-github-bot commented Aug 4, 2021 •

edited

Loading

codecov Bot commented Sep 16, 2021 •

edited

Loading

mruberry Feb 22, 2022 •

edited

Loading