Add subdiagonalCount:superdiagonalCount: argument labels to Tensor.bandPart. (#588)

Argument labels improve clarity.
Deprecate label-less version, which has one user in tensorflow/swift-models.

Author: dan-zheng

Sources/TensorFlow/Operators/LinearAlgebra.swift

@@ -56,6 +56,12 @@ public extension Tensor where Scalar: TensorFlowNumeric {
         _Raw.matrixDiag(diagonal: self)
     }
 
+    @available(*, deprecated, renamed: "bandPart(subdiagonalCount:superdiagonalCount:)")
+    @differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
+    func bandPart(_ subdiagonalCount: Int, _ superdiagonalCount: Int) -> Tensor {
+        return bandPart(subdiagonalCount: subdiagonalCount, superdiagonalCount: superdiagonalCount)
+    }
+
     /// Returns a copy of a innermost tensor defined by a central band boundaries.
     /// The output is a tensor of the same shape as the instance `[..., :, :]`.
     ///
@@ -79,12 +85,18 @@ public extension Tensor where Scalar: TensorFlowNumeric {
     /// //  [-2, -1,  0, 1]
     /// //  [ 0, -2, -1, 0]]
     /// ```
+    ///
+    /// - Parameters:
+    ///   - subdiagonalCount: The number of subdiagonals to keep. If negative, keep entire lower
+    ///     triangle.
+    ///   - superdiagonalCount: The number of superdiagonals to keep. If negative, keep entire upper
+    ///     triangle.
     @inlinable
     @differentiable(wrt: self, vjp: _vjpBandPart where Scalar: TensorFlowFloatingPoint)
-    func bandPart(_ lowerCount: Int, _ upperCount: Int) -> Tensor {
+    func bandPart(subdiagonalCount: Int, superdiagonalCount: Int) -> Tensor {
         precondition(rank >= 2, "The tensor must have at least rank 2.")
-        let lower = Tensor<Int32>(Int32(lowerCount))
-        let upper = Tensor<Int32>(Int32(upperCount))
+        let lower = Tensor<Int32>(Int32(subdiagonalCount))
+        let upper = Tensor<Int32>(Int32(superdiagonalCount))
         return _Raw.matrixBandPart(self, numLower: lower, numUpper: upper)
     }
 }
@@ -101,8 +113,15 @@ internal extension Tensor where Scalar: TensorFlowFloatingPoint {
     }
 
     @inlinable
-    func _vjpBandPart(_ numLower: Int, _ numUpper: Int) -> (Tensor, (Tensor) -> Tensor) {
-        (bandPart(numLower, numUpper), { $0.bandPart(numLower, numUpper) })
+    func _vjpBandPart(
+        subdiagonalCount: Int, superdiagonalCount: Int
+    ) -> (Tensor, (Tensor) -> Tensor) {
+        let value = bandPart(
+            subdiagonalCount: subdiagonalCount,
+            superdiagonalCount: superdiagonalCount)
+        return (value, {
+            $0.bandPart(subdiagonalCount: subdiagonalCount, superdiagonalCount: superdiagonalCount)
+        })
     }
 }
 

Tests/TensorFlowTests/OperatorTests/MatrixTests.swift

@@ -80,16 +80,17 @@ final class MatrixTests: XCTestCase {
                                      [-1,  0,  1, 2],
                                      [ 0, -1,  0, 1],
                                      [ 0,  0, -1, 0]])
-        XCTAssertEqual(t1.bandPart(1, -1), target1)
+        XCTAssertEqual(t1.bandPart(subdiagonalCount: 1, superdiagonalCount: -1), target1)
 
         let target2 = Tensor<Float>([[ 0,  1,  0, 0],
                                      [-1,  0,  1, 0],
                                      [-2, -1,  0, 1],
                                      [ 0, -2, -1, 0]])
-        XCTAssertEqual(t1.bandPart(2, 1), target2)
+        XCTAssertEqual(t1.bandPart(subdiagonalCount: 2, superdiagonalCount: 1), target2)
 
         // Test special case - diagonal
-        XCTAssertEqual(t1.bandPart(0, 0), Tensor<Float>(zeros: [4, 4]))
+        XCTAssertEqual(t1.bandPart(subdiagonalCount: 0, superdiagonalCount: 0),
+                       Tensor<Float>(zeros: [4, 4]))
 
         // Test leading dimensions with special case - lower triangular
         let t2 = Tensor<Float>(stacking: [t1, t1 + 1])
@@ -101,12 +102,14 @@ final class MatrixTests: XCTestCase {
                                       [ 0,  1,  0, 0],
                                       [-1,  0,  1, 0],
                                       [-2, -1,  0, 1]]])
-        XCTAssertEqual(t2.bandPart(-1, 0), target3)
+        XCTAssertEqual(t2.bandPart(subdiagonalCount: -1, superdiagonalCount: 0), target3)
 
         // Test bandPart gradient with special case - upper triangular
         let t3 = Tensor<Float>(shape: [2, 4, 4], scalars: (1...(2 * 16)).map(Float.init))
-        let computedGrad = gradient(at: t3) { $0.squared().bandPart(0, -1).sum() }
-        let expectedGrad = 2 * t3.bandPart(0, -1)
+        let computedGrad = gradient(at: t3) {
+            $0.squared().bandPart(subdiagonalCount: 0, superdiagonalCount: -1).sum()
+        }
+        let expectedGrad = 2 * t3.bandPart(subdiagonalCount: 0, superdiagonalCount: -1)
         XCTAssertEqual(computedGrad, expectedGrad)
     }