Add topk xla op.

Author: pschuh

Sources/CX10/xla_tensor_wrapper.cc

@@ -651,6 +651,11 @@ OpaqueXLATensor* XLATensor_tan(OpaqueXLATensor* a) {
 OpaqueXLATensor* XLATensor_tanh(OpaqueXLATensor* a) {
   return new XLATensor(XLATensor::tanh(*a));
 }
+OpaqueXLATensor_pair XLATensor_topk(OpaqueXLATensor* a, int64_t k,
+                                            int64_t dim, bool largest) {
+  auto result = XLATensor::topk(*a, k, dim, largest, false);
+  return {new XLATensor(std::get<0>(result)), new XLATensor(std::get<1>(result))};
+}
 OpaqueXLATensor* XLATensor_tf_Conv(OpaqueXLATensor* input,
                                    OpaqueXLATensor* filter, bool depthwise,
                                    Int64ArrayRef strides, TFPadding padding,

Sources/CX10/xla_tensor_wrapper.h

@@ -374,6 +374,8 @@ XLA_API OpaqueXLATensor* XLATensor_sum(OpaqueXLATensor* a, Int64ArrayRef dims,
                                        Optional_XLAScalarType dtype);
 XLA_API OpaqueXLATensor* XLATensor_tan(OpaqueXLATensor* a);
 XLA_API OpaqueXLATensor* XLATensor_tanh(OpaqueXLATensor* a);
+XLA_API OpaqueXLATensor_pair XLATensor_topk(OpaqueXLATensor* a, int64_t k,
+                                            int64_t dim, bool largest);
 XLA_API OpaqueXLATensor*
 XLATensor_tf_Conv(OpaqueXLATensor* input, OpaqueXLATensor* filter, bool depthwise,
                   Int64ArrayRef strides, enum TFPadding padding,

Sources/x10/swift_bindings/XLATensor.swift

@@ -407,7 +407,9 @@ extension XLATensor {
     return XLATensor(_handle: XLATensor_div(a.handle, b.handle))
   }
 
-  static func dynamic_slice(_ base: XLATensor, _ start_indices: [XLATensor], _ slice_shape: [Int64]) -> XLATensor {
+  static func dynamic_slice(_ base: XLATensor, _ start_indices: [XLATensor], _ slice_shape: [Int64])
+    -> XLATensor
+  {
     start_indices.withArrayRef { start_indices in
       slice_shape.withArrayRef { slice_shape in
         return XLATensor(_handle: XLATensor_dynamic_slice(base.handle, start_indices, slice_shape))
@@ -761,7 +763,7 @@ extension XLATensor {
   }
 
   static func replica_id(_ device: Device) -> XLATensor {
-    return XLATensor(_handle: XLATensor_replica_id(device.cdevice));
+    return XLATensor(_handle: XLATensor_replica_id(device.cdevice))
   }
 
   static func resize_value(_ value: XLATensor, _ dims: [Int64]) -> XLATensor {
@@ -874,6 +876,12 @@ extension XLATensor {
     return XLATensor(_handle: XLATensor_tanh(a.handle))
   }
 
+  static func topk(_ a: XLATensor, k: Int64, dim: Int64, largest: Bool) -> (XLATensor, XLATensor) {
+    defer { _fixLifetime(a) }
+    let output = XLATensor_topk(a.handle, k, dim, largest)
+    return (XLATensor(_handle: output.x), XLATensor(_handle: output.y))
+  }
+
   static func tf_Conv(
     _ input: XLATensor, _ filter: XLATensor, _ depthwise: Bool, _ strides: [Int64],
     _ padding: TFPadding, _ explicit_paddings: [Int64],

Sources/x10/swift_bindings/apis/RawOpsManual.swift

@@ -1425,7 +1425,8 @@ public enum _RawXLA {
     _ startIndices: [Tensor<Int32>],
     _ sliceShape: [Int64]
   ) -> Tensor<T> {
-    return Tensor(_xla: XLATensor.dynamic_slice(base.xlaTensor, startIndices.map { $0.xlaTensor }, sliceShape))
+    return Tensor(
+      _xla: XLATensor.dynamic_slice(base.xlaTensor, startIndices.map { $0.xlaTensor }, sliceShape))
   }
 
   public static func dynamicUpdateSlice<T: TensorFlowNumeric>(
@@ -3042,17 +3043,17 @@ public enum _RawXLA {
   ///
   /// NOTE `tf.reverse` has now changed behavior in preparation for 1.0.
   /// `tf.reverse_v2` is currently an alias that will be deprecated before TF 1.0.
-  /// 
+  ///
   /// Given a `tensor`, and a `int32` tensor `axis` representing the set of
   /// dimensions of `tensor` to reverse. This operation reverses each dimension
   /// `i` for which there exists `j` s.t. `axis[j] == i`.
-  /// 
+  ///
   /// `tensor` can have up to 8 dimensions. The number of dimensions specified
   /// in `axis` may be 0 or more entries. If an index is specified more than
   /// once, a InvalidArgument error is raised.
-  /// 
+  ///
   /// For example:
-  /// 
+  ///
   /// ```
   /// # tensor 't' is [[[[ 0,  1,  2,  3],
   /// #                  [ 4,  5,  6,  7],
@@ -3061,23 +3062,23 @@ public enum _RawXLA {
   /// #                  [16, 17, 18, 19],
   /// #                  [20, 21, 22, 23]]]]
   /// # tensor 't' shape is [1, 2, 3, 4]
-  /// 
+  ///
   /// # 'dims' is [3] or 'dims' is [-1]
   /// reverse(t, dims) ==> [[[[ 3,  2,  1,  0],
   ///                         [ 7,  6,  5,  4],
   ///                         [ 11, 10, 9, 8]],
   ///                        [[15, 14, 13, 12],
   ///                         [19, 18, 17, 16],
   ///                         [23, 22, 21, 20]]]]
-  /// 
+  ///
   /// # 'dims' is '[1]' (or 'dims' is '[-3]')
   /// reverse(t, dims) ==> [[[[12, 13, 14, 15],
   ///                         [16, 17, 18, 19],
   ///                         [20, 21, 22, 23]
   ///                        [[ 0,  1,  2,  3],
   ///                         [ 4,  5,  6,  7],
   ///                         [ 8,  9, 10, 11]]]]
-  /// 
+  ///
   /// # 'dims' is '[2]' (or 'dims' is '[-2]')
   /// reverse(t, dims) ==> [[[[8, 9, 10, 11],
   ///                         [4, 5, 6, 7],
@@ -4259,6 +4260,13 @@ public enum _RawXLA {
     return Tensor(_xla: XLATensor.tanh(x.xlaTensor))
   }
 
+  public static func topk<T: FloatingPoint & TensorFlowScalar>(
+    _ a: Tensor<T>, k: Int64, dim: Int64, largest: Bool
+  ) -> (Tensor<T>, Tensor<Int64>) {
+    let (r0, r1) = XLATensor.topk(a.xlaTensor, k: k, dim: dim, largest: largest)
+    return (Tensor(_xla: r0), Tensor(_xla: r1))
+  }
+
   /// Assign `value` to the sliced l-value reference of `input`.
   ///
   /// The values of `value` are assigned to the positions in the tensor `input` that