cuda_tensor.norm(dim=(X, Y)) is broken

[ssnl]

In [1]: torch.__version__
Out[1]: '1.4.0a0+91c6d2e'

In [2]: torch.randn(3, 3, device='cuda').norm(dim=(0, 1))
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-2-32861c8633ea> in <module>
----> 1 torch.randn(3, 3, device='cuda').norm(dim=(0, 1))

/data/packages/pytorch/torch/tensor.py in norm(self, p, dim, keepdim, dtype)
    337     def norm(self, p="fro", dim=None, keepdim=False, dtype=None):
    338         r"""See :func:`torch.norm`"""
--> 339         return torch.norm(self, p, dim, keepdim, dtype=dtype)
    340
    341     def lu(self, pivot=True, get_infos=False):

/data/packages/pytorch/torch/functional.py in norm(input, p, dim, keepdim, out, dtype)
    780             dim = tuple(range(ndim))
    781         if out is None:
--> 782             return torch._C._VariableFunctions.frobenius_norm(input, dim, keepdim=keepdim)
    783         return torch._C._VariableFunctions.frobenius_norm(input, dim, keepdim=keepdim, out=out)
    784     elif p == "nuc":

RuntimeError: Could not run 'aten::conj.out' with arguments from the 'CUDATensorId' backend. 'aten::conj.out' is only available for these backends: [CPUTensorId, Va
riableTensorId].

In [3]: torch.randn(3, 3, device='cuda').norm()
Out[3]: tensor(3.6015, device='cuda:0')

In [4]: torch.randn(3, 3).norm(dim=(0, 1))
Out[4]: tensor(3.5681)

In [5]: torch.randn(3, 3).norm()
Out[5]: tensor(2.7136)

Yet passing in a 2-tuple is perfectly valid, per the torch.norm doc:

dim (int, 2-tuple of python:ints, 2-list of python:ints, optional) – If it is an int, vector norm will be calculated, if it is 2-tuple of ints, matrix norm will be calculated. If the value is None, matrix norm will be calculated when the input tensor only has two dimensions, vector norm will be calculated when the input tensor only has one dimension. If the input tensor has more than two dimensions, the vector norm will be applied to last dimension.

cc @ezyang @gchanan @zou3519 @anjali411 @dylanbespalko

[zou3519]

Based on the documentation, this sounds like a regression so I am marking it as high-pri. We should figure out if it is actually a regression (I haven't tested this on previous versions of pytorch).

[ngimel]

This is likely a regression resulting from #27653

[ezyang]

cc @dylanbespalko

[ezyang]

It sounds like the correct fix is to make sure conj (and real) work on CUDA tensors as well as CPU tensors. Actually, it's possible that #30295 just up and fixes this.

[ezyang]

Also @dylanbespalko, it seems like the modified version of norm and similar sites are a pessimization for the non-complex case, because it looks like your implementations of conj/real always do "work" (in the sense that they have to allocate and fill a new tensor). So regardless it may be a good idea to avoid calling these functions at all for non-complex inputs.

[dylanbespalko]

@ezyang,

Yes, this was definitely caused by me. This should effect the frobenius norm in aten/src/ATen/native/LinearAlgebra.cpp. The line of code in question is

at::sqrt(at::sum((self.conj() * self).real(), dim, keepdim));

And can be fixed two ways:

1. Adding CUDA support for Unary ops in CUDA-strided-complex Binary and Unary Op support #30295.
2. Creating a separate code-path for non-complex and complex cases.
   • @ezyang points out that conj() and real() are not in-place operations, so fix 1 will still have performance degradation.

Solution:
Fix 1 should resolve the bug, but I will also submit a PR of fix 2 to improve the performance.