---

yaml
---
r: 312286
b: refs/heads/tensorflow-merge
c: b30deb2
h: refs/heads/master

Author: rxwei

[refs]

@@ -1379,7 +1379,7 @@ refs/heads/chase-my-tail: 8bb91443a9e81bbfac92a2621a0af887a1da8dbf
 refs/heads/consider-outer-alternatives: 708bac749ec60a22a79e2eefbe734f9488a7370d
 refs/heads/revert-25740-oops-i-linked-it-again: fdd41aeb682fc488572bdc1cf71b2ff6997ba576
 refs/heads/swift-5.1-branch-06-12-2019: e63b7b2d3b93c48232d386099d0ec525d21d8f8d
-refs/heads/tensorflow-merge: 3e8a1deb09ebecd6c93d44b02874a8a3c89fe5dc
+refs/heads/tensorflow-merge: b30deb2d3d1720ac4e0e2041fa9f4c63d51fcb3b
 refs/heads/update-checkout-sha-info: 5832743c5c2a842976c42a508a4c6dcceefb0aef
 refs/tags/swift-5.1-DEVELOPMENT-SNAPSHOT-2019-06-12-a: 228f0448d9bb909aacbba4afcb7c600a405d15da
 refs/tags/swift-5.1-DEVELOPMENT-SNAPSHOT-2019-06-14-a: 922861a77b5fc2bf46bc917da70ceb15eef76836

branches/tensorflow-merge/stdlib/public/TensorFlow/Gradients.swift

@@ -569,25 +569,15 @@ extension Tensor where Scalar : TensorFlowFloatingPoint {
 
 extension Tensor where Scalar : TensorFlowFloatingPoint {
   @inlinable
-  func _vjpSum() -> (Tensor, (Tensor) -> Tensor) {
-    return (sum(), { [shape = shapeTensor] in $0.broadcast(toShape: shape) })
-  }
-
-  @inlinable
-  func _vjpMean() -> (Tensor, (Tensor) -> Tensor) {
-    return (mean(), { [shape = shapeTensor, count = scalarCountTensor] in
-      ($0 / Tensor(count)).broadcast(toShape: shape)
-    })
-  }
-
-  @inlinable
-  func _vjpSum(alongAxes axes: [Int32]) -> (Tensor, (Tensor) -> Tensor) {
+  func _vjpSum(alongAxes axes: Tensor<Int32>) -> (Tensor, (Tensor) -> Tensor) {
     let value = sum(alongAxes: axes)
     return (value, { [shape = shapeTensor] in $0.broadcast(toShape: shape) })
   }
 
   @inlinable
-  func _vjpSum(squeezingAxes axes: [Int32]) -> (Tensor, (Tensor) -> Tensor) {
+  func _vjpSum(
+    squeezingAxes axes: Tensor<Int32>
+  ) -> (Tensor, (Tensor) -> Tensor) {
     let value = sum(squeezingAxes: axes)
     return (value, { [shape = shapeTensor] in $0.broadcast(toShape: shape) })
   }
@@ -602,20 +592,13 @@ extension Tensor where Scalar : TensorFlowFloatingPoint {
   }
 
   @inlinable
-  func _vjpMean(squeezingAxes axes: [Int32]) -> (Tensor, (Tensor) -> Tensor) {
+  func _vjpMean(
+    squeezingAxes axes: Tensor<Int32>
+  ) -> (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(alongAxes axes: [Int32]) -> (Tensor, (Tensor) -> Tensor) {
-    let value = mean(alongAxes: axes)
-    return (value, { [shape = shapeTensor,
-                      count = axes.map { shape[$0] }.reduce(1, *)] in
-      $0.broadcast(toShape: shape) / Tensor(Scalar(count))
+    let count = Raw.gather(params: shapeTensor, indices: axes).product()
+    return (value, { [shape = shapeTensor] in
+      $0.broadcast(toShape: shape) / Tensor(count)
     })
   }
 }