Add torch.linalg.trace

pytorch · asi1024 · Aug 4, 2021 · Sep 13, 2021 · Sep 21, 2021 · Sep 29, 2021
commit 0b91b4b4183525dea9dec7a60636df7231c9b7c9
diff --git a/aten/src/ATen/native/ReduceOps.cpp b/aten/src/ATen/native/ReduceOps.cpp
@@ -1079,6 +1079,19 @@ Tensor trace_cpu(const Tensor& self) {
   return result;
 }
 
+// TODO: this routine should be implemented without diag and sum for perf problems,
+// see https://github.com/pytorch/pytorch/pull/47305,
+Tensor linalg_trace(const Tensor& self, int64_t offset) {
+  return at::diagonal(self, offset, -2, -1).sum(-1);
+}
+
+Tensor linalg_trace_backward(const Tensor & grad, IntArrayRef input_sizes, int64_t offset) {
+  auto grad_input = at::zeros(input_sizes, grad.options());
+  auto diag = grad_input.diagonal(offset, -2, -1);
+  diag.copy_(grad.unsqueeze(-1));
+  return grad_input;
+}
+
 void impl_func_prod(
     const Tensor& self,
     IntArrayRef dims,
@@ -1093,6 +1106,13 @@ void impl_func_prod(
   }
 }
 
+Tensor prod(const Tensor& self, int64_t dim, bool keepdim, c10::optional<ScalarType> opt_dtype) {
 structured_delegate: prod.int_out 
 structured_delegate: prod.int_out 
+  ScalarType dtype = get_dtype_from_self(self, opt_dtype, true);
+  Tensor result = create_reduction_result(self, dim, keepdim, dtype);
+  native::prod_out_impl(result, self, dim, keepdim, dtype);
+  return result;
+}
+
 TORCH_IMPL_FUNC(prod_out)
 (const Tensor& self,
  int64_t dim,

diff --git a/aten/src/ATen/native/native_functions.yaml b/aten/src/ATen/native/native_functions.yaml
@@ -6230,6 +6230,18 @@
     CPU: trace_cpu
     CUDA: trace_cuda
 
+- func: linalg_trace(Tensor self, int offset=0) -> Tensor
+  python_module: linalg
+  variants: method, function
+  dispatch:
+    CPU, CUDA: linalg_trace
-    CPU, CUDA: linalg_trace
-    CPU, CUDA: linalg_trace
+    CompositeExplicitAutograd: linalg_trace
+
+- func: linalg_trace_backward(Tensor grad, int[] sizes, int offset) -> Tensor
+  variants: function
+  device_check: NoCheck
+  device_guard: False
+
 - func: trace_backward(Tensor grad, int[] sizes) -> Tensor
   variants: function
   device_check: NoCheck

diff --git a/tools/autograd/derivatives.yaml b/tools/autograd/derivatives.yaml
@@ -1412,6 +1412,9 @@
   self: trace_backward(grad, self.sizes())
   result: auto_linear
 
+- name: linalg_trace(Tensor self, int offset=0) -> Tensor
+  self: linalg_trace_backward(grad, self.sizes(), offset)
+
 - name: transpose.int(Tensor(a) self, int dim0, int dim1) -> Tensor(a)
   self: grad.transpose(dim0, dim1)
   result: auto_linear

diff --git a/torch/linalg/__init__.py b/torch/linalg/__init__.py
@@ -2042,3 +2042,19 @@
     >>> torch.dist(Q.transpose(-2, -1) @ Q, torch.eye(4))
     tensor(6.2158e-07)
 """)
+
+trace = _add_docstr(_linalg.linalg_trace, r"""
+trace(input, offset=0) -> Tensor
+
+Returns the sum of the elements of the diagonal.
+
+Example::
+
+    >>> x = torch.arange(1., 10.).view(3, 3)
+    >>> x
+    tensor([[ 1.,  2.,  3.],
+            [ 4.,  5.,  6.],
+            [ 7.,  8.,  9.]])
+    >>> torch.linalg.trace(x)
+    tensor(15.)
+""")
diff --git a/torch/overrides.py b/torch/overrides.py
@@ -958,6 +958,7 @@ def get_testing_overrides() -> Dict[Callable, Callable]:
         torch.tile: lambda input, dims: -1,
         torch.topk: lambda input, k, dim=-1, descending=False, out=None: -1,
         torch.trace: lambda input: -1,
+        torch.linalg.trace: lambda input, offset=0: -1,
         torch.transpose: lambda input, dim0, dim1: -1,
         torch.trapz: lambda y, x=None, dim=-1: -1,
         torch.trapezoid: lambda y, x=None, dim=-1: -1,