[TF] Added support for advanced indexing and slicing

stdlib/public/TensorFlow/Ops.swift

@@ -1598,38 +1598,6 @@ public extension Tensor where Scalar : Numeric {
 //===----------------------------------------------------------------------===//
 
 public extension Tensor {
-  /// Access the element tensor specified by an index in the leading dimension.
-  /// - Parameter index: Index of the element tensor.
-  @inlinable
-  subscript(index: Int32) -> Tensor {
-    get {
-      let slice = Raw.stridedSlice(
-        self, begin: Tensor<Int32>([index]), end: Tensor<Int32>([index + 1]),
-        strides: Tensor<Int32>([1]))
-      return slice.squeezingShape(at: 0)
-    }
-    set {
-      let leftElements = self[0..<index]
-      let rightElements = self[index+1..<shape[0]]
-      self = Raw.concatV2([leftElements, newValue.rankLifted(), rightElements],
-                          axis: Tensor<Int32>(0))
-    }
-  }
-
-  /// Access the subtensor specified by a contiguous range of indices.
-  /// - Parameter bounds: Contiguous range of indices.
-  @inlinable
-  subscript(bounds: Range<Int32>) -> Tensor {
-    return Raw.stridedSlice(
-      self, begin: Tensor<Int32>([bounds.lowerBound]),
-      end: Tensor<Int32>([bounds.upperBound]), strides: Tensor<Int32>([1]))
-  }
-
-  // TODO(danielzheng): Add strided slices? (increment by something different
-  // than 1)
-  // Ideas for strided slice API: it could be another subscript method, or it
-  // be a top level `stride` function like Swift's `stride(from:to:by:)`.
-
   /// Extracts a slice from the tensor defined by lower and upper bounds for
   /// each dimension.
   ///
@@ -1646,3 +1614,165 @@ public extension Tensor {
       size: Tensor<Int32>(upperBounds) - lowerBoundsTensor)
   }
 }
+
+public enum TensorSliceIndex : TensorSliceIndexProtocol {
+  case ellipsis
+  case newAxis
+  case squeezeAxis
+  case index(Int32)
+  case range(Range<Int32>, stride: Int32)
+  case closedRange(ClosedRange<Int32>, stride: Int32)
+  case partialRangeFrom(PartialRangeFrom<Int32>, stride: Int32)
+  case partialRangeUpTo(PartialRangeUpTo<Int32>, stride: Int32)
+  case partialRangeThrough(PartialRangeThrough<Int32>, stride: Int32)
+
+  public var sliceIndex: TensorSliceIndex { return self }
+}
+
+public protocol TensorSliceIndexProtocol {
+  var sliceIndex: TensorSliceIndex { get }
+}
+
+// TODO: Cannot extend non-nominal type 'UnboundedRange'.
+// extension UnboundedRange : TensorSliceIndexProtocol {
+//   public var sliceIndex: TensorSliceIndex { return .ellipsis }
+// }
+
+extension Int32 : TensorSliceIndexProtocol {
+  public var sliceIndex: TensorSliceIndex { return .index(self) }
+}
+
+extension Int : TensorSliceIndexProtocol {
+  public var sliceIndex: TensorSliceIndex { return .index(Int32(self)) }
+}
+
+extension Range : TensorSliceIndexProtocol where Bound == Int {
+  public var sliceIndex: TensorSliceIndex {
+    return .range(Int32(self.lowerBound)..<Int32(self.upperBound), stride: 1)
+  }
+}
+
+extension ClosedRange : TensorSliceIndexProtocol where Bound == Int {
+  public var sliceIndex: TensorSliceIndex {
+    return .closedRange(Int32(self.lowerBound)...Int32(self.upperBound), stride: 1)
+  }
+}
+
+extension PartialRangeFrom : TensorSliceIndexProtocol where Bound == Int {
+  public var sliceIndex: TensorSliceIndex {
+    return .partialRangeFrom(Int32(self.lowerBound)..., stride: 1)
+  }
+}
+
+extension PartialRangeUpTo : TensorSliceIndexProtocol where Bound == Int {
+  public var sliceIndex: TensorSliceIndex {
+    return .partialRangeUpTo(..<Int32(self.upperBound), stride: 1)
+  }
+}
+
+extension PartialRangeThrough : TensorSliceIndexProtocol where Bound == Int {
+  public var sliceIndex: TensorSliceIndex {
+    return .partialRangeThrough(...Int32(self.upperBound), stride: 1)
+  }
+}
+
+public extension Tensor {
+  public struct IndexPath {
+    public let begin, end, strides: Tensor<Int32>
+    public let beginMask, endMask, ellipsisMask, newAxisMask, squeezeAxisMask: Int64
+  }
+
+  @inlinable @inline(__always)
+  subscript(_ indexPath: IndexPath) -> Tensor {
+    get {
+      return Raw.stridedSlice(
+        self, begin: indexPath.begin, end: indexPath.end, strides: indexPath.strides,
+        beginMask: indexPath.beginMask, endMask: indexPath.endMask, 
+        ellipsisMask: indexPath.ellipsisMask, newAxisMask: indexPath.newAxisMask, 
+        shrinkAxisMask: indexPath.squeezeAxisMask)
+    }
+    set {
+      Raw.tensorStridedSliceUpdate(
+        self, begin: indexPath.begin, end: indexPath.end, strides: indexPath.strides,
+        value: newValue, beginMask: indexPath.beginMask, endMask: indexPath.endMask, 
+        ellipsisMask: indexPath.ellipsisMask, newAxisMask: indexPath.newAxisMask, 
+        shrinkAxisMask: indexPath.squeezeAxisMask)
+    }
+  }
+
+  @inlinable @inline(__always)
+  subscript(_ indices: TensorSliceIndexProtocol...) -> Tensor {
+    get {
+      return self[IndexPath(indices.map { $0.sliceIndex })]
+    }
+    set {
+      self[IndexPath(indices.map { $0.sliceIndex })] = newValue
+    }
+  }
+}
+
+public extension Tensor.IndexPath {
+  @inlinable @inline(__always)
+  init(_ indices: [TensorSliceIndex]) {
+    precondition(!indices.isEmpty, "The index path cannot be empty.")
+    precondition(indices.count(where: {
+      if case .ellipsis = $0 {
+        return true
+      } else {
+        return false
+      }
+    }) < 2, "Only one ellipsis is allowed per index path.")
+
+    var begin = [Int32](repeating: 0, count: indices.count)
+    var end = [Int32](repeating: 0, count: indices.count)
+    var strides = [Int32](repeating: 1, count: indices.count)
+    var beginMask: Int64 = 0
+    var endMask: Int64 = 0
+    var ellipsisMask: Int64 = 0
+    var newAxisMask: Int64 = 0
+    var squeezeAxisMask: Int64 = 0
+    for (i, index) in indices.enumerated() {
+      switch index {
+      case .ellipsis: ellipsisMask |= 1 << i
+      case .newAxis: newAxisMask |= 1 << i
+      case .squeezeAxis: squeezeAxisMask |= 1 << i
+      case .index(let idx):
+        begin[i] = idx
+        end[i] = idx + 1
+        squeezeAxisMask |= 1 << i
+      case .range(let range, let stride):
+        begin[i] = range.lowerBound
+        end[i] = range.upperBound
+        strides[i] = stride
+      case .closedRange(let range, let stride):
+        begin[i] = range.lowerBound
+        switch range.upperBound {
+        case -1: endMask |= 1 << i
+        case let u: end[i] = u + 1
+        }
+        strides[i] = stride
+      case .partialRangeFrom(let range, let stride):
+        begin[i] = range.lowerBound
+        strides[i] = stride
+        endMask |= 1 << i
+      case .partialRangeUpTo(let range, let stride):
+        end[i] = range.upperBound
+        strides[i] = stride
+        beginMask |= 1 << i
+      case .partialRangeThrough(let range, let stride):
+        end[i] = range.upperBound + 1
+        strides[i] = stride
+        beginMask |= 1 << i
+      }
+    }
+
+    self.begin = Tensor<Int32>(begin)
+    self.end = Tensor<Int32>(end)
+    self.strides = Tensor<Int32>(strides)
+    self.beginMask = beginMask
+    self.endMask = endMask
+    self.ellipsisMask = ellipsisMask
+    self.newAxisMask = newAxisMask
+    self.squeezeAxisMask = squeezeAxisMask
+  }
+}

test/TensorFlowRuntime/tensor.swift

@@ -135,6 +135,29 @@ TensorTests.testAllBackends("ElementIndexing") {
   expectEqual([43], array0D.scalars)
 }
 
+TensorTests.testAllBackends("NestedElementIndexing") {
+  // NOTE: This tests the `subscript(indices:)` method, which is distinct from
+  // the `subscript(index:)` method.
+  // NOTE: This test could use a clearer name, along with other "indexing"
+  // tests. Note to update corresponding test names in other files
+  // (shaped_array.test) as well.
+  let tensor3D = Tensor<Float>(shape: [3, 4, 5],
+                               scalars: Array(stride(from: 0.0, to: 60, by: 1)))
+  let element1D = tensor3D[1, 3]
+  let element0D = tensor3D[2, 0, 3]
+
+  let array1D = element1D.array
+  let array0D = element0D.array
+
+  /// Test shapes
+  expectEqual([5], array1D.shape)
+  expectEqual([], array0D.shape)
+
+  /// Test scalars
+  expectEqual(Array(stride(from: 35.0, to: 40, by: 1)), array1D.scalars)
+  expectEqual([43], array0D.scalars)
+}
+
 TensorTests.testAllBackends("SliceIndexing") {
   // XLA compilation error under TPU.
   if _RuntimeConfig.executionMode.isTPU { return }
@@ -173,6 +196,23 @@ TensorTests.test("WholeTensorSlicing") {
               slice2.array)
 }
 
+TensorTests.testAllBackends("AdvancedIndexing") {
+  // NOTE: cannot test multiple `Tensor.shape` or `Tensor.scalars` directly
+  // until send and receive are implemented (without writing a bunch of mini
+  // tests). Instead, `Tensor.array` is called to make a ShapedArray host copy
+  // and the ShapedArray is tested.
+  let tensor3D = Tensor<Float>(shape: [3, 4, 5],
+                               scalars: Array(stride(from: 0.0, to: 60, by: 1)))
+  let element2D = tensor3D[1 ..< 3, 0, 3...]
+  let array2D = element2D.array
+
+  // Test shape
+  expectEqual([2, 2], array2D.shape)
+
+  // Test scalars
+  expectEqual(Array([23.0, 24.0, 43.0, 44.0]), array2D.scalars)
+}
+
 TensorTests.testAllBackends("Reduction") {
   // TODO(b/111815968): triage and fix this TPU issue
   #if !TPU

utils/build-script-impl

@@ -2444,10 +2444,11 @@ for host in "${ALL_HOSTS[@]}"; do
                     # problematic for overwriting/stripping symbols. Thus, write
                     # permission is added here.
                     for lib_name in tensorflow tensorflow_framework; do
-                        lib="lib${lib_name}.so"
+                        lib="lib${lib_name}.so*"
+                        dylib="lib${lib_name}*.dylib"
                         rm -f "${TF_LIB_DIR}/${lib}"
-                        chmod +w "${TENSORFLOW_HOST_LIB_DIR}/${lib}"
-                        cp -p "${TENSORFLOW_HOST_LIB_DIR}/${lib}" "${TF_LIB_DIR}"
+                        find "${TENSORFLOW_HOST_LIB_DIR}" \( -name "${lib}" -o -name "${dylib}" \) -exec chmod +w {} +
+                        find "${TENSORFLOW_HOST_LIB_DIR}" \( -name "${lib}" -o -name "${dylib}" \) -exec cp -p {} "${TF_LIB_DIR}" \;
                     done
 
                     if [[ ! "${TENSORFLOW_TARGET_LIB_DIR}" ]] ; then
@@ -3828,10 +3829,11 @@ for host in "${ALL_HOSTS[@]}"; do
                 mkdir -p "${TF_DEST_DIR}"
                 for lib_name in tensorflow tensorflow_framework
                 do
-                    lib="lib${lib_name}.so"
+                    lib="lib${lib_name}.so*"
+                    dylib="lib${lib_name}*.dylib"
                     echo "${TF_LIBDIR}/${lib} => ${TF_DEST_DIR}"
                     rm -f "${TF_DEST_DIR}/${lib}"
-                    cp -a "${TF_LIBDIR}/${lib}" "${TF_DEST_DIR}"
+                    find "${TF_LIBDIR}" \( -name "${lib}" -o -name "${dylib}" \) -exec cp -a {} "${TF_DEST_DIR}" \;
                 done
                 continue
                 ;;

utils/update_checkout/update-checkout-config.json

@@ -229,7 +229,7 @@
                 "lldb": "bdb96e8b352f7cb18e2b3b66f2d3f75d92f81dcd",
                 "cmark": "swift-DEVELOPMENT-SNAPSHOT-2018-11-26-a",
                 "llbuild": "swift-DEVELOPMENT-SNAPSHOT-2018-11-26-a",
-                "swiftpm": "swift-DEVELOPMENT-SNAPSHOT-2018-11-26-a",
+                "swiftpm": "f7deeaac96dc42475e15d045bb093436c98f89ee",
                 "swift-syntax": "1aeb642da66a23a66c9ac80d74813e1a4b963999",
                 "swift-stress-tester": "2fc093642df924f6adf9de9e4397c7c6fc8b5fc8",
                 "compiler-rt": "swift-DEVELOPMENT-SNAPSHOT-2018-11-26-a",
@@ -240,7 +240,7 @@
                 "swift-xcode-playground-support": "swift-DEVELOPMENT-SNAPSHOT-2018-11-26-a",
                 "ninja": "253e94c1fa511704baeb61cf69995bbf09ba435e",
                 "icu": "release-61-1",
-                "tensorflow": "5e8df789cc30098d791475c14a623ec68b50b4ed",
+                "tensorflow": "5f54f6497068f25dd574ad0a45b6ab2e2e920561",
                 "tensorflow-swift-bindings": "0957744551614e433dbabc725cba29ff5ddb91d3",
                 "tensorflow-swift-apis": "5caa4600e0796cc04dc755f7d7c4befe7bd336cd"
             }