Revert non-active try-apply differentiation workarounds (#1069)

Since TF-433 was fixed, we can remove all the workarkounds for (non-active) try-apply.
For example, calls to `Array.map(_:)`.

Author: efremale

Sources/TensorFlow/Operators/Math.swift

@@ -1860,8 +1860,7 @@ extension Tensor where Scalar: Numeric & Comparable {
   @inlinable
   @differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
   public func max(squeezingAxes axes: [Int]) -> Tensor {
-    // TODO(TF-433): Remove workaround for differentiating `map`.
-    let axes = { axes.map(Int32.init) }()
+    let axes = axes.map(Int32.init)
     return max(squeezingAxes: Tensor<Int32>(axes, on: device))
   }
 
@@ -1890,8 +1889,7 @@ extension Tensor where Scalar: Numeric & Comparable {
   @inlinable
   @differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
   public func min(squeezingAxes axes: [Int]) -> Tensor {
-    // TODO(TF-433): Remove workaround for differentiating `map`.
-    let axes = { axes.map(Int32.init) }()
+    let axes = axes.map(Int32.init)
     return min(squeezingAxes: Tensor<Int32>(axes, on: device))
   }
 
@@ -1942,8 +1940,7 @@ extension Tensor where Scalar: Numeric & Comparable {
   @inlinable
   @differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
   public func min(alongAxes axes: [Int]) -> Tensor {
-    // TODO(TF-433): Remove workaround for differentiating `map`.
-    let axes = { axes.map(Int32.init) }()
+    let axes = axes.map(Int32.init)
     return min(alongAxes: Tensor<Int32>(axes, on: device))
   }
 
@@ -1975,8 +1972,7 @@ extension Tensor where Scalar: Numeric & Comparable {
   @inlinable
   @differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
   public func max(alongAxes axes: [Int]) -> Tensor {
-    // TODO(TF-433): Remove workaround for differentiating `map`.
-    let axes = { axes.map(Int32.init) }()
+    let axes = axes.map(Int32.init)
     return max(alongAxes: Tensor<Int32>(axes, on: device))
   }
 
@@ -2114,8 +2110,7 @@ extension Tensor where Scalar: Numeric {
   @inlinable
   @differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
   public func sum(squeezingAxes axes: [Int]) -> Tensor {
-    // TODO(TF-433): Remove workaround for differentiating `map`.
-    let axes = { axes.map(Int64.init) }()
+    let axes = axes.map(Int64.init)
     return _Raw.sum(self, reductionIndices: axes, keepDims: false)
   }
 
@@ -2150,8 +2145,7 @@ extension Tensor where Scalar: Numeric {
   @inlinable
   @differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
   public func sum(alongAxes axes: [Int]) -> Tensor {
-    // TODO(TF-433): Remove workaround for differentiating `map`.
-    let axes = { axes.map(Int64.init) }()
+    let axes = axes.map(Int64.init)
     return _Raw.sum(self, reductionIndices: axes, keepDims: true)
   }
 
@@ -2184,8 +2178,7 @@ extension Tensor where Scalar: Numeric {
   @inlinable
   @differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
   public func product(squeezingAxes axes: [Int]) -> Tensor {
-    // TODO(TF-433): Remove workaround for differentiating `map`.
-    let axes = { axes.map(Int32.init) }()
+    let axes = axes.map(Int32.init)
     return product(squeezingAxes: Tensor<Int32>(axes, on: device))
   }
 
@@ -2221,8 +2214,7 @@ extension Tensor where Scalar: Numeric {
   /// - Precondition: Each value in `axes` must be in the range `-rank..<rank`.
   @inlinable
   public func product(alongAxes axes: [Int]) -> Tensor {
-    // TODO(TF-433): Remove workaround for differentiating `map`.
-    let axes = { axes.map(Int32.init) }()
+    let axes = axes.map(Int32.init)
     return product(alongAxes: Tensor<Int32>(axes, on: device))
   }
 
@@ -2253,8 +2245,7 @@ extension Tensor where Scalar: Numeric {
   @inlinable
   @differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
   public func mean(squeezingAxes axes: [Int]) -> Tensor {
-    // TODO(TF-433): Remove workaround for differentiating `map`.
-    let axes = { axes.map(Int64.init) }()
+    let axes = axes.map(Int64.init)
     return _Raw.mean(self, reductionIndices: axes, keepDims: false)
   }
 
@@ -2291,8 +2282,7 @@ extension Tensor where Scalar: Numeric {
   @inlinable
   @differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
   public func mean(alongAxes axes: [Int]) -> Tensor {
-    // TODO(TF-433): Remove workaround for differentiating `map`.
-    let axes = { axes.map(Int64.init) }()
+    let axes = axes.map(Int64.init)
     return _Raw.mean(self, reductionIndices: axes, keepDims: true)
   }
 
@@ -2327,8 +2317,7 @@ extension Tensor where Scalar: Numeric {
   @inlinable
   @differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
   public func variance(squeezingAxes axes: [Int]) -> Tensor {
-    // TODO(TF-433): Remove workaround for differentiating `map`.
-    let axes = { axes.map(Int32.init) }()
+    let axes = axes.map(Int32.init)
     return variance(squeezingAxes: Tensor<Int32>(axes, on: device))
   }
 
@@ -2369,8 +2358,7 @@ extension Tensor where Scalar: Numeric {
   @inlinable
   @differentiable(wrt: self where Scalar: TensorFlowFloatingPoint)
   public func variance(alongAxes axes: [Int]) -> Tensor {
-    // TODO(TF-433): Remove workaround for differentiating `map`.
-    let axes = { axes.map(Int32.init) }()
+    let axes = axes.map(Int32.init)
     return variance(alongAxes: Tensor<Int32>(axes, on: device))
   }
 
@@ -2791,8 +2779,7 @@ extension Tensor where Scalar: TensorFlowFloatingPoint {
   @inlinable
   @differentiable(wrt: self)
   public func standardDeviation(alongAxes axes: [Int]) -> Tensor {
-    // TODO(TF-433): Remove workaround for differentiating `map`.
-    let axes = { axes.map(Int32.init) }()
+    let axes = axes.map(Int32.init)
     return standardDeviation(alongAxes: Tensor<Int32>(axes, on: device))
   }
 
@@ -2842,8 +2829,7 @@ extension Tensor where Scalar: TensorFlowFloatingPoint {
   @inlinable
   @differentiable(wrt: self)
   public func logSumExp(squeezingAxes axes: [Int]) -> Tensor {
-    // TODO(TF-433): Remove workaround for differentiating `map`.
-    let axes = withoutDerivative(at: axes) { $0.map(Int32.init) }
+    let axes = axes.map(Int32.init)
     return logSumExp(squeezingAxes: Tensor<Int32>(axes, on: device))
   }
 
@@ -2907,8 +2893,7 @@ extension Tensor where Scalar: TensorFlowFloatingPoint {
   @inlinable
   @differentiable(wrt: self)
   public func logSumExp(alongAxes axes: [Int]) -> Tensor {
-    // TODO(TF-433): Remove workaround for differentiating `map`.
-    let axes = withoutDerivative(at: axes) { $0.map(Int32.init) }
+    let axes = axes.map(Int32.init)
     return logSumExp(alongAxes: Tensor<Int32>(axes, on: device))
   }