---

yaml
---
r: 312314
b: refs/heads/tensorflow-merge
c: dc6415e
h: refs/heads/master

Author: sgugger

[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: 1ab9f04028722d7803b777b859bb4db83ead6a42
+refs/heads/tensorflow-merge: dc6415ef0ea6e918c26bc04f4d89cea7a9bc1546
 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

@@ -553,12 +553,10 @@ extension Tensor where Scalar : TensorFlowFloatingPoint {
   }
 
   @inlinable
-  func _vjpExpandingShape(
-    at shapeIndex: Int
-  ) -> (Tensor, (Tensor) -> Tensor) {
-    let value = expandingShape(at: shapeIndex)
+  func _vjpExpandingShape(at axes: [Int]) -> (Tensor, (Tensor) -> Tensor) {
+    let value = self.expandingShape(at: axes)
     return (value, { v in
-      v.squeezingShape(at: shapeIndex)
+      v.squeezingShape(at: axes)
     })
   }
 }

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

@@ -706,14 +706,24 @@ public extension Tensor {
   }
 
   /// Returns a shape-expanded `Tensor`, with a dimension of 1 inserted at the
-  /// specified shape index.
+  /// specified shape indices.
+  @inlinable @inline(__always)
+  @differentiable(wrt: self where Scalar : TensorFlowFloatingPoint)
+  func expandingShape(at axes: Int...) -> Tensor {
+    return expandingShape(at: axes)
+  }
+   
+  /// Returns a shape-expanded `Tensor`, with a dimension of 1 inserted at the
+  /// specified shape indices.
   @inlinable @inline(__always)
   @differentiable(
     wrt: self, vjp: _vjpExpandingShape(at:)
     where Scalar : TensorFlowFloatingPoint
   )
-  func expandingShape(at shapeIndex: Int) -> Tensor {
-    return Raw.expandDims(self, dim: Tensor<Int32>(Int32(shapeIndex)))
+  func expandingShape(at axes: [Int]) -> Tensor {
+    var res = self
+    for i in axes { res = Raw.expandDims(res, dim: Tensor<Int32>(Int32(i))) }
+    return res
   }
 
   /// Remove the specified dimensions of size 1 from the shape of a tensor. If

branches/tensorflow-merge/test/TensorFlowRuntime/tensor.swift

@@ -674,6 +674,20 @@ TensorTests.testAllBackends("MLPClassifierStruct") {
   expectPointwiseNearlyEqual([0.816997], prediction.scalars)
 }
 
+TensorTests.testAllBackends("ExpandingShape") {
+  // 2 x 3 -> 1 x 2 x 1 x 3 x 1
+  let matrix = Tensor<Int32>([[0, 1, 2], [3, 4, 5]])
+  let reshaped = matrix.expandingShape(at: 0,2,4)
+
+  expectEqual([1, 2, 1, 3, 1], reshaped.shape)
+  expectEqual(Array(0..<6), reshaped.scalars)
+  
+  // 1 x 2 x 1 x 3 x 1 -> 2 x 3
+  let rereshaped = reshaped.squeezingShape(at: 0,2,4)
+  expectEqual([2, 3], rereshaped.shape)
+  expectEqual(Array(0..<6), rereshaped.scalars)
+}
+
 TensorTests.testAllBackends("Reshape") {
   // 2 x 3 -> 1 x 3 x 1 x 2 x 1
   let matrix = Tensor<Int32>([[0, 1, 2], [3, 4, 5]])