Makes 'Tensor.gathering' generic over 'Int32' and 'Int64'. (#272)

Added a 'TensorFlowIndex' protocol and make 'Int32' and 'Int64' conform to it.

Author: eaplatanios

Sources/TensorFlow/Core/DataTypes.swift

@@ -70,6 +70,13 @@ public typealias TensorFlowNumeric = TensorFlowScalar & Numeric
 public typealias TensorFlowSignedNumeric = TensorFlowScalar & SignedNumeric
 public typealias TensorFlowInteger = TensorFlowScalar & BinaryInteger
 
+/// An integer data type that represents integer types which can be used as tensor indices in 
+/// TensorFlow.
+public protocol TensorFlowIndex: TensorFlowInteger {}
+
+extension Int32: TensorFlowIndex {}
+extension Int64: TensorFlowIndex {}
+
 /// A floating-point data type that conforms to `Differentiable` and is compatible with TensorFlow.
 ///
 /// - Note: `Tensor` conditionally conforms to `Differentiable` when the `Scalar` associated type

Sources/TensorFlow/Operators/Basic.swift

@@ -383,7 +383,10 @@ public extension Tensor {
     /// - Returns: The gathered tensor.
     @inlinable
     @differentiable(wrt: self, vjp: _vjpGathering where Scalar : TensorFlowFloatingPoint)
-    func gathering(atIndices indices: Tensor<Int32>, alongAxis axis: Int = 0) -> Tensor {
+    func gathering<Index: TensorFlowIndex>(
+        atIndices indices: Tensor<Index>,
+        alongAxis axis: Int = 0
+    ) -> Tensor {
         return Raw.gatherV2(params: self, indices: indices, axis: Tensor<Int32>(Int32(axis)))
     }
 
@@ -408,15 +411,15 @@ public extension Tensor {
     /// - Returns: The gathered tensor.
     @inlinable
     @differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
-    func batchGathering(atIndices indices: Tensor<Int32>) -> Tensor {
+    func batchGathering<Index: TensorFlowIndex>(atIndices indices: Tensor<Index>) -> Tensor {
         var batchIndices = indices
-        var accumulated = Tensor<Int32>(ones: [])
-        accumulated *= Swift.withoutDerivative(at: shapeTensor) { $0[1] }
+        var accumulated = Tensor<Index>(ones: [])
+        accumulated *= Swift.withoutDerivative(at: shapeTensor) { Tensor<Index>($0[1]) }
         let dValue = Swift.withoutDerivative(at: shapeTensor) { $0[0] }
-        let dIndices = Tensor<Int32>(
-            rangeFrom: Tensor<Int32>(zeros: []),
-            to: dValue,
-            stride: Tensor<Int32>(ones: [])
+        let dIndices = Tensor<Index>(
+            rangeFrom: Tensor<Index>(zeros: []),
+            to: Tensor<Index>(dValue),
+            stride: Tensor<Index>(ones: [])
         ) * accumulated
         let dShape = Tensor<Int32>(concatenating: [
             dValue.rankLifted(),
@@ -519,8 +522,8 @@ internal extension Tensor where Scalar: TensorFlowFloatingPoint {
     }
 
     @inlinable
-    func _vjpGathering(
-        atIndices indices: Tensor<Int32>,
+    func _vjpGathering<Index: TensorFlowIndex>(
+        atIndices indices: Tensor<Index>,
         alongAxis axis: Int = 0
     ) -> (Tensor, (Tensor) -> Tensor) {
         let result = gathering(atIndices: indices, alongAxis: axis)