Refine tiled (#689)

Add tiled API and tests with refined preconditions.

Sources/TensorFlow/Operators/Basic.swift

@@ -132,6 +132,25 @@ public extension Tensor {
             splitDim: Tensor<Int32>(Int32(axis)),
             numSplit: Int64(sizes.shape[0]))
     }
+
+    /// Returns a tiled tensor, constructed by tiling this tensor.
+    ///
+    /// This constructor creates a new tensor by replicating this tensor `multiples` times. The
+    /// constructed tensor's `i`'th dimension has `self.shape[i] * multiples[i]` elements, and the
+    /// values of this tensor are replicated `multiples[i]` times along the `i`'th dimension. For
+    /// example, tiling `[a b c d]` by `[2]` produces `[a b c d a b c d]`.
+    ///
+    /// - Precondition: The expected `rank` of multiples must be `1`.
+    /// - Precondition: The shape of `multiples` must be `[tensor.rank]`.
+    /// - Precondition: All scalars in `multiples` must be non-negative.
+    @inlinable
+    @differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
+    func tiled(multiples: [Int]) -> Tensor {
+        precondition(multiples.allSatisfy { $0 >= 0 },
+                     "All scalars in multiples must be non-negative.")
+        // TODO(TF-433): Remove workaround for differentiating `map`.
+        return tiled(multiples: Tensor<Int32>({multiples.map(Int32.init)}()))
+    }
 
     /// Returns a tiled tensor, constructed by tiling this tensor.
     ///

Tests/TensorFlowTests/OperatorTests/BasicTests.swift

@@ -475,6 +475,16 @@ final class BasicOperatorTests: XCTestCase {
         XCTAssertEqual(grad.shape, [3, 2, 1])
         XCTAssertEqual(grad.scalars, [1, 1, 1, 1, 1, 1])
     }
+
+    func testTile() {
+        let tensor = Tensor<Int32>([[0, 1, 2], [3, 4, 5]])
+        let tiled = tensor.tiled(multiples: [3, 2])
+
+        XCTAssertEqual(tiled.shape, [6, 6])
+        XCTAssertEqual(tiled, [[0, 1, 2, 0, 1, 2], [3, 4, 5, 3, 4, 5],
+                               [0, 1, 2, 0, 1, 2], [3, 4, 5, 3, 4, 5],
+                               [0, 1, 2, 0, 1, 2], [3, 4, 5, 3, 4, 5]])
+    }
 
     func testReshape() {
         // 2 x 3 -> 1 x 3 x 1 x 2 x 1
@@ -692,6 +702,7 @@ final class BasicOperatorTests: XCTestCase {
         ("testConcatenation", testConcatenation),
         ("testVJPConcatenation", testVJPConcatenation),
         ("testTranspose", testTranspose),
+        ("testTile", testTile),
         ("testReshape", testReshape),
         ("testFlatten", testFlatten),
         ("testFlatten0D", testFlatten0D),

Tests/TensorFlowTests/TensorAutoDiffTests.swift

@@ -482,6 +482,15 @@ final class TensorAutoDiffTests: XCTestCase {
         XCTAssertEqual(pullback(at: [[3, 5]], in: f1)([1, 1]), [[6, 10]])
         XCTAssertEqual(pullback(at: [[3, 5]], in: f2)([1, 1]), [[6, 10]])
     }
+
+    func testTiled() {
+        let input = Tensor<Float>([[1, 2, 3], [4, 5, 6]])
+        let tiledPullback = pullback(at: input) { (a: Tensor<Float>) in
+            a.tiled(multiples: [2, 1])
+        }
+        let tiled = Tensor<Float>([[1, 2, 3], [4, 5, 6], [1, 2, 3], [4, 5, 6]])
+        XCTAssertEqual(input * 2, tiledPullback(tiled))
+    }
 
     func testReshapedBackprop() {
         func f1(a: Tensor<Float>) -> Tensor<Float> { a.reshaped(toShape: Tensor<Int32>([2, 1])).squared() }
@@ -795,6 +804,7 @@ final class TensorAutoDiffTests: XCTestCase {
         ("testTensorInitStacking", testTensorInitStacking),
         ("testExpandingShape", testExpandingShape),
         ("testSqueezingShape", testSqueezingShape),
+        ("testTiled", testTiled),
         ("testReshapedBackprop", testReshapedBackprop),
         ("testReshaped", testReshaped),
         ("testConcatenationPlusPlus", testConcatenationPlusPlus),