Made 'Tensor.gathering(where:alongAxis:)' differentiable. (#271)

Author: eaplatanios

Sources/TensorFlow/Operators/Basic.swift

@@ -464,12 +464,10 @@ public extension Tensor {
     /// - Returns: `(self.rank - K + 1)`-dimensional tensor populated by entries in this tensor
     ///   corresponding to `true` values in `mask`.
     @inlinable
-    // @differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
+    @differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
     func gathering(where mask: Tensor<Bool>, alongAxis axis: Int = 0) -> Tensor {
         precondition(mask.rank != 0, "The boolean mask cannot be a scalar.")
-        // TODO: Remove once control flow AD is supported.
-        let rank = self.rank
-        let posAxis = { axis < 0 ? axis + rank : axis }()
+        let posAxis = Swift.withoutDerivative(at: self.rank) { r in axis < 0 ? axis + r : axis }
         let leadingSize = shapeTensor[posAxis ..< posAxis + mask.rank].product().rankLifted()
         let reshapedTensor = reshaped(
             toShape: Tensor<Int32>(concatenating: [