[Linear Algebra] Trace operator (#586)

Adds a trace operator and adds tests and derivative tests.

Author: awav

Sources/TensorFlow/Operators/LinearAlgebra.swift

@@ -106,6 +106,21 @@ internal extension Tensor where Scalar: TensorFlowFloatingPoint {
     }
 }
 
+/// Computes the trace of an optionally batched matrix.
+/// The trace is the the sum along the main diagonal of each inner-most matrix.
+///
+/// The input is a tensor with shape `[..., M, N]`.
+/// The output is a tensor with shape `[...]`.
+///
+/// - Parameter matrix: A tensor of shape `[..., M, N]`.
+/// - Precondition: `matrix` must be a tensor with shape `[..., M, N]`.
+@inlinable
+@differentiable(wrt: matrix where T: TensorFlowFloatingPoint)
+public func trace<T: TensorFlowNumeric>(_ matrix: Tensor<T>) -> Tensor<T> {
+    precondition(matrix.rank >= 2, "The tensor must have at least rank 2.")
+    return matrix.diagonalPart().sum(squeezingAxes: -1)
+}
+
 // MARK: - Decompositions
 
 /// Returns the Cholesky decomposition of one or more square matrices.

Tests/TensorFlowTests/OperatorTests/LinearAlgebraTests.swift

@@ -60,8 +60,35 @@ final class LinearAlgebraTests: XCTestCase {
         }
     }
 
+    func testTrace() {
+        assertEqual(trace(Tensor<Float>(ones: [3, 3])), Tensor(3.0), accuracy: 1e-16)
+        assertEqual(trace(Tensor<Float>(ones: [5, 6])), Tensor(5.0), accuracy: 1e-16)
+        let shapes = [[1, 3, 3], [2, 4, 4], [2, 3, 5, 5]]
+        for shape in shapes {
+            let x = Tensor<Float>(ones: TensorShape(shape))
+            let computedTrace = trace(x)
+            let leadingShape = TensorShape(shape.dropLast(2))
+            let value = Float(shape.last!)
+            let expectedTrace = Tensor<Float>(repeating: value, shape: leadingShape)
+            assertEqual(computedTrace, expectedTrace, accuracy: 1e-16)
+        }
+    }
+
+    func testTraceGradient() {
+        let shape: TensorShape = [2, 4, 4]
+        let scalars = (0..<shape.contiguousSize).map(Float.init)
+        let x = Tensor<Float>(shape: shape, scalars: scalars)
+        let computedGradient = gradient(at: x) { (trace($0) * [2.0, 3.0]).sum() }
+        let a = Tensor<Float>(repeating: 2.0, shape: [4]).diagonal()
+        let b = Tensor<Float>(repeating: 3.0, shape: [4]).diagonal()
+        let expectedGradient = Tensor([a, b])
+        assertEqual(computedGradient, expectedGradient, accuracy: 1e-16)
+    }
+    
     static var allTests = [
         ("testCholesky", testCholesky),
-        ("testQRDecompositionApproximation", testQRDecompositionApproximation)
+        ("testQRDecompositionApproximation", testQRDecompositionApproximation),
+        ("testTrace", testTrace),
+        ("testTraceGradient", testTraceGradient),
     ]
 }