Fixed a bug in 'Tensor.batchGathering(atIndices:alongAxis:batchDimensionsCount:)'. (#359)

Author: eaplatanios

Sources/TensorFlow/Operators/Basic.swift

@@ -414,12 +414,15 @@ public extension Tensor {
         alongAxis axis: Int = 1,
         batchDimensionCount: Int = 1
     ) -> Tensor {
-        // TODO: precondition(batchDimensionCount >= 0,
-        //                    "'batchDimensionCount' must be non-negative.")
-        // TODO: precondition(batchDimensionCount < indices.rank,
-        //                    "'batchDimensionCount' must be less than 'indices.rank'.")
-        // TODO: precondition(batchDimensionCount < rank, 
-        //                    "'batchDimensionCount' must be less than the tensor's rank.")
+        precondition(batchDimensionCount >= 0, "'batchDimensionCount' must be non-negative.")
+        precondition(
+            batchDimensionCount < indices.rank,
+            "'batchDimensionCount' must be less than 'indices.rank'.")
+        withoutDerivative(at: rank) {
+            precondition(
+                batchDimensionCount < $0,
+                "'batchDimensionCount' must be less than the tensor's rank.")
+        }
 
         // Handle the axis argument by transposing the axis dimension so that it is the first
         // non-batch dimension, recursively calling `batchGathering` with `axis = 0`, and then
@@ -472,7 +475,7 @@ public extension Tensor {
                 let dShape = Tensor<Int32>(concatenating: [
                     Tensor<Int32>([Int32](repeating: 1, count: d - 1)),
                     dValue.rankLifted(),
-                    Tensor<Int32>([Int32](repeating: 1, count: indices.rank - 1))])
+                    Tensor<Int32>([Int32](repeating: 1, count: indices.rank - d))])
                 batchIndices += dIndices.reshaped(toShape: dShape)
             }
             return batchIndices

Tests/TensorFlowTests/OperatorTests/BasicTests.swift

@@ -27,10 +27,10 @@ final class BasicOperatorTests: XCTestCase {
             [1.0, 2.0, 3.0],
             [4.0, 5.0, 6.0]]])
         let y = x.batchGathering(
-            atIndices: Tensor<Int32>([1, 0]),
+            atIndices: Tensor<Int32>([[[1], [0]]]),
             alongAxis: 2,
             batchDimensionCount: 2)
-        XCTAssertEqual(y, Tensor<Float>([2.0, 4.0]))
+        XCTAssertEqual(y, Tensor<Float>([[[2.0], [4.0]]]))
     }
 
     func testPadded() {