[TF] Fix gradients in sum(squeezinAxes:) and mean(squeezinAxes:) (#24164)

* [TF] Fix gradients in sum(squeezinAxes:) and mean(squeezinAxes:)

sum(squeezinAxes:) and mean(squeezinAxes:) were throwing an error during the bawckward pass because the gradients weren't unsqueezed before being broadcast.
Note that this could be refactored nicely if we had a function that took a list of ints for `expandingShape`.
Second note: I may be wrong, but it seems like `_vjpMean(squeezingAxes axes: [Int])` is never used and only the Tensor<Int32> version is.

* Remove unused `_vjpMean` function.

* Update Gradients.swift

* Add test

* Minor edit for consistency.

Author: sgugger

stdlib/public/TensorFlow/Gradients.swift

@@ -579,7 +579,11 @@ extension Tensor where Scalar : TensorFlowFloatingPoint {
     squeezingAxes axes: Tensor<Int32>
   ) -> (Tensor, (Tensor) -> Tensor) {
     let value = sum(squeezingAxes: axes)
-    return (value, { [shape = shapeTensor] in $0.broadcast(toShape: shape) })
+    return (value, { [shape = shapeTensor] in
+      var res = $0
+      for i in axes.array.scalars { res = res.expandingShape(at: Int(i)) }
+      return res.broadcast(toShape: shape)
+    })
   }
 
   @inlinable
@@ -591,23 +595,16 @@ extension Tensor where Scalar : TensorFlowFloatingPoint {
     })
   }
 
-  @inlinable
-  func _vjpMean(squeezingAxes axes: [Int]) -> (Tensor, (Tensor) -> Tensor) {
-    let value = mean(squeezingAxes: axes)
-    return (value, { [shape = shapeTensor,
-                      count = axes.map { shape[$0] }.reduce(1, *)] in
-      $0.broadcast(toShape: shape) / Tensor(Scalar(count))
-    })
-  }
-
   @inlinable
   func _vjpMean(
     squeezingAxes axes: Tensor<Int32>
   ) -> (Tensor, (Tensor) -> Tensor) {
     let value = mean(squeezingAxes: axes)
     let count = Raw.gather(params: shapeTensor, indices: axes).product()
     return (value, { [shape = shapeTensor] in
-      $0.broadcast(toShape: shape) / Tensor(count)
+      var res = $0
+      for i in axes.array.scalars { res = res.expandingShape(at: Int(i)) }
+      return res.broadcast(toShape: shape) / Tensor(count)
     })
   }
 }

test/TensorFlowRuntime/tensor_autodiff_runtime.swift

@@ -98,38 +98,39 @@ TensorADTests.testAllBackends("Abs") {
 TensorADTests.testAllBackends("sum") {
   let input = Tensor<Float>(repeating: 42, shape: [2, 2])
   let sumPullbackScalar = pullback(at: input) { (a: Tensor<Float>) in a.sum() }
+  let sumPullbackSqueezingAxes = pullback(at: input) { (a: Tensor<Float>) in a.sum(squeezingAxes: 0, 1) }
   let sumPullbackAlongAxes = pullback(at: input) { (a: Tensor<Float>) in a.sum(alongAxes: 0, 1) }
 
   let expected = Tensor<Float>(ones: [2, 2])
   expectEqual(expected, sumPullbackScalar(Tensor(1)))
-  // expectEqual(expected, sumPullbackSqueezingAxes(Tensor(1)))
+  expectEqual(expected, sumPullbackSqueezingAxes(Tensor(1)))
   expectEqual(expected, sumPullbackAlongAxes(Tensor(1)))
   expectEqual(expected * 3, sumPullbackScalar(Tensor(3)))
-  // expectEqual(expected * 3, sumPullbackSqueezingAxes(Tensor(3)))
+  expectEqual(expected * 3, sumPullbackSqueezingAxes(Tensor(3)))
   expectEqual(expected * 3, sumPullbackAlongAxes(Tensor(3)))
 }
 
 TensorADTests.testAllBackends("mean") {
   let meanGradScalar = gradient { (a: Tensor<Float>) in a.mean() }
-  // let meanGradSqueezingAxes = gradient { (a: Tensor<Float>) in a.mean(squeezingAxes: 0, 1) }
+  let meanGradSqueezingAxes = gradient { (a: Tensor<Float>) in a.mean(squeezingAxes: 0, 1) }
   let meanGradAlongAxes = gradient { (a: Tensor<Float>) in a.mean(alongAxes: 0, 1) }
 
   let input = Tensor<Float>(ones: [2, 2])
   let expected = Tensor<Float>(repeating: 0.25, shape: [2, 2])
   expectEqual(expected, meanGradScalar(input))
-  // expectEqual(expected, meanGradSqueezingAxes(input))
+  expectEqual(expected, meanGradSqueezingAxes(input))
   expectEqual(expected, meanGradAlongAxes(input))
 }
 
 TensorADTests.testAllBackends("variance") {
   let varianceGradScalar = gradient { (a: Tensor<Float>) in a.variance() }
-  // let varianceGradSqueezingAxes = gradient { (a: Tensor<Float>) in a.variance(squeezingAxes: 0, 1) }
+  let varianceGradSqueezingAxes = gradient { (a: Tensor<Float>) in a.variance(squeezingAxes: 0, 1) }
   let varianceGradAlongAxes = gradient { (a: Tensor<Float>) in a.variance(alongAxes: 0, 1) }
 
   let input: Tensor<Float> = [[1, 2], [3, 4]]
   let expected: Tensor<Float> = [[-0.75, -0.25], [0.25, 0.75]]
   expectEqual(expected, varianceGradScalar(input))
-  // expectEqual(expected, varianceGradSqueezingAxes(input))
+  expectEqual(expected, varianceGradSqueezingAxes(input))
   expectEqual(expected, varianceGradAlongAxes(input))
 }