Add missing complex support for torch.norm and torch.linalg.norm #48284

kurtamohler · 2020-11-20T00:21:57Z

BC-breaking note:

Previously, when given a complex input, torch.linalg.norm and torch.norm would return a complex output. torch.linalg.cond would sometimes return a complex output and sometimes return a real output when given a complex input, depending on its p argument. This PR changes this behavior to match numpy.linalg.norm and numpy.linalg.cond, so that a complex input will result in the downgraded real number type, consistent with NumPy.

PR Summary:

The following cases were previously unsupported for complex inputs, and this commit adds support:

Frobenius norm
Norm order 2 (vector and matrix)
CUDA vector norm

Part of #47833

dr-ci · 2020-11-20T00:33:42Z

💊 CI failures summary and remediations

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

💚 💚 Looks good so far! There are no failures yet. 💚 💚

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

Follow this link to opt-out of these comments for your Pull Requests.

Please report bugs/suggestions on the GitHub issue tracker or post in the (internal) Dr. CI Users group.

See how this bot performed.

This comment has been revised 64 times.

test/test_linalg.py

kurtamohler · 2020-11-24T18:54:32Z

Rocm test failure seems to be real

aten/src/ATen/native/SharedReduceOps.h

kurtamohler · 2020-11-24T21:20:42Z

@mruberry, in numpy, if we give norm a complex number, it returns the downgraded real type:

>>> numpy.linalg.norm([1+1j]).dtype
dtype('float64')

But in pytorch, we match the input type:

>>> torch.linalg.norm(torch.tensor([1+1j])).dtype
torch.complex64

Is it alright if we BC-break to fix this?

codecov · 2020-11-24T21:23:11Z

Codecov Report

Merging #48284 (df1244f) into master (274ce26) will decrease coverage by 0.00%.
The diff coverage is 93.33%.

@@            Coverage Diff             @@
##           master   #48284      +/-   ##
==========================================
- Coverage   80.74%   80.73%   -0.01%     
==========================================
  Files        1869     1869              
  Lines      201654   201650       -4     
==========================================
- Hits       162818   162802      -16     
- Misses      38836    38848      +12

aten/src/ATen/native/ReduceOps.cpp

aten/src/ATen/native/SharedReduceOps.h

mruberry · 2020-11-27T02:25:13Z

aten/src/ATen/native/cpu/ReduceOpsKernel.cpp

@@ -174,61 +174,72 @@ static void norm_kernel_tensor_iterator_impl(
  if (p.isIntegral(false)) {
    val = p.to<int64_t>();
  } else if (p.isFloatingPoint()) {
-    val = p.to<float>();
+    val = p.to<double>();


Why the change from float to double here?

To me, it seems better to cast to double to allow more precision for p.

aten/src/ATen/native/SharedReduceOps.h

mruberry · 2020-11-27T02:33:28Z

aten/src/ATen/native/SharedReduceOps.h

-  inline C10_DEVICE acc_t reduce(acc_t acc, acc_t data, int64_t /*idx*/) const {
-    return acc + data * data;
+  inline C10_DEVICE acc_t reduce(acc_t acc, scalar_t data, int64_t /*idx*/) const {
+    acc_t abs_data = std::abs(data);


What's the difference between the previous expression and the new one with std::abs?

The std::abs() call is needed when scalar_t is complex. For real numbers, it's not needed. I could add an overload for complex numbers so that we avoid calling std::abs() for real numbers.

Adding an overload sounds good as long as it doesn't add too much code complexity.

I added abs_if_complex(). Not too sure if this is the most efficient way to implement it. Let me know if you think this adds too much complexity.

aten/src/ATen/native/cpu/ReduceOpsKernel.cpp

aten/src/ATen/native/cuda/ReduceNormKernel.cu

mruberry · 2020-11-27T02:44:16Z

Hey @kurtamohler!

This PR looks great. It has a surgical precision. I made a few comments that are mostly me asking questions to better understand what's going on.

I think this PR should update the documentation, too, to more accurately describe torch.norm's and torch.linalg.norm's complex support.

Also, how are we testing complex autograd for torch.linalg.norm? Should something like this be updated?

https://github.com/pytorch/pytorch/blob/bb5d4984b912f3f9f775b00b31de198fd3d01a7f/test/test_linalg.py#L937

Also, the lint build should be fixed if you rebase.

kurtamohler · 2020-11-30T19:14:39Z

Thanks for the review @mruberry! Yes, more accurate documentation is a good idea, and I will look into autograd testing

torch/linalg/__init__.py

mruberry

Awesome! Thanks @kurtamohler.

Two small cleanup suggestions in the docs and one question about using torch.promote_types() to simplify a test. Just let me know when this is ready to merge.

kurtamohler · 2020-12-04T17:40:42Z

@mruberry, I think this is ready to merge

kurtamohler · 2020-12-04T20:22:51Z

Looks like the return dtype change (return real when given complex) is breaking one of the tests for torch.linalg.cond because it depends on torch.linalg.norm. In numpy, cond has the same return type behavior as norm:

>>> a = numpy.random.rand(10, 10) + 1j * numpy.random.rand(10, 10)
>>> a.dtype
dtype('complex128')
>>> numpy.linalg.cond(a).dtype
dtype('float64')

So I can fix the failing errors and add cond to the BC breaking note above.

kurtamohler · 2020-12-05T01:23:29Z

I'm not sure what's causing the pytorch-linux-bionic-rocm3.9-py3.6 failure. I'll look into it

mruberry · 2020-12-06T03:03:37Z

I'm not sure what's causing the pytorch-linux-bionic-rocm3.9-py3.6 failure. I'll look into it

It doesn't look related. We can probably ignore it (unless the same failure happens again).

kurtamohler · 2020-12-07T18:28:49Z

In that case, I think this is ready to merge @mruberry , unless another test fails after I rebased

kurtamohler · 2020-12-07T21:34:16Z

All the failures are due to an upstream mypy issue

aten/src/ATen/native/ReduceOps.cpp

torch/functional.py

mruberry · 2020-12-08T09:34:42Z

torch/linalg/__init__.py

+        will be returned. Its data type must be either a floating point or complex type. For complex
+        inputs, the norm is calculated on of the absolute values of each element. If the input is
+        complex and neither :attr:`dtype` nor :attr:`out` is specified, the result's data type will be
+        the corresponding downgraded real number type.


This sentence needs to be updated to be consistent with the torch/functional.py documentation.

Does the torch.linalg.cond documentation need a similar update, too?

Updated this and torch.linalg.cond

mruberry · 2020-12-08T09:37:55Z

Hey @kurtamohler! Just a few comments/questions. Overall things look great.

facebook-github-bot

@mruberry has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator.

facebook-github-bot · 2020-12-10T19:17:54Z

@mruberry merged this pull request in 54f0556.

anjali411 · 2021-01-04T19:02:27Z

aten/src/ATen/native/LinearAlgebra.cpp

@@ -1679,7 +1677,7 @@ static Tensor& _linalg_norm_vector_out(Tensor& result, const Tensor& self, optio
      // when the input contains extreme values (like nan or +/-inf) or if the input
      // size is degenerate (like size(0), size(0, N), etc)
      case_was_overridden = true;
-      self_ = self.abs();
+      self_ = self_.abs();


why don't we just declare self_ without assignment above?

The conditional dtype conversion opt_dtype.has_value() ? self.to(opt_dtype.value()) : self needs to be performed for all cases, which is the reason why self_ is defined

facebook-github-bot added the cla signed label Nov 20, 2020

pytorchbot added the open source label Nov 20, 2020

kurtamohler force-pushed the norm-complex-support-47833 branch 2 times, most recently from ba0ae51 to 1408afd Compare November 24, 2020 00:08

kurtamohler commented Nov 24, 2020

View reviewed changes

test/test_linalg.py Outdated Show resolved Hide resolved

kurtamohler force-pushed the norm-complex-support-47833 branch from 1408afd to d279241 Compare November 24, 2020 18:09

kurtamohler changed the title ~~WIP: Add missing complex support for torch.norm and torch.linalg.norm~~ Add missing complex support for torch.norm and torch.linalg.norm Nov 24, 2020

kurtamohler force-pushed the norm-complex-support-47833 branch from d279241 to 684bbd6 Compare November 24, 2020 18:14

kurtamohler marked this pull request as ready for review November 24, 2020 18:15

kurtamohler requested a review from mruberry November 24, 2020 18:15

kurtamohler commented Nov 24, 2020

View reviewed changes

aten/src/ATen/native/SharedReduceOps.h Outdated Show resolved Hide resolved

kurtamohler linked an issue Nov 24, 2020 that may be closed by this pull request

Support torch.linalg.norm for complex tensors on both CPU and CUDA #47833

Closed

kurtamohler force-pushed the norm-complex-support-47833 branch 2 times, most recently from d80710e to bb5d498 Compare November 25, 2020 22:14

ngimel added the triaged This issue has been looked at a team member, and triaged and prioritized into an appropriate module label Nov 26, 2020