XlaCompileOnDemandOp);                               \
  REGISTER_KERNEL_BUILDER(                                                     \
      Name("XlaVariadicSort").HostMemory("dimension").Device(DEVICE),          \
      XlaCompileOnDemandOp);                                                   \
  REGISTER_KERNEL_BUILDER(Name("XlaWhile").Device(DEVICE),                     \

                                         "XlaSpmdShardToFullShape",
                                         "XlaSvd",
                                         "XlaVariadicReduceV2",
                                         "XlaVariadicSort",
                                         "XlaWhile"};

static bool NodeCanTriggerXlaCompilation(const NodeDef& node) {
  return node.attr().find(kXlaClusterIdAttr) != node.attr().end() ||

      "XlaSort",
      "XlaSplitND",
      "XlaSpmdFullToShardShape",
      "XlaSpmdShardToFullShape",
      "XlaSvd",
      "XlaVariadicReduce",
      "XlaVariadicReduceV2",
      "XlaVariadicSort",
      "XlaWhile",
      "Zeta",
      "_Arg",
      "_ArrayToList",
      "_ListToArray",
      "_Retval"};
  return result;
}

}  // namespace testing

  %0 = "tf.AddV2"(%arg1, %arg0) : (tensor<i32>, tensor<i32>) -> tensor<i32>
  func.return %0 : tensor<i32>
}

//===----------------------------------------------------------------------===//
// tf.XlaVariadicSort legalization
//===----------------------------------------------------------------------===//

// CHECK-LABEL: @xla_variadic_sort
// CHECK-SAME: %[[INPUT:.*]]: tensor<2x3x4xui8>

    BuildBodyWithCall(rewriter, loc, func, func_ty, &reduce_op.getBody());

    rewriter.replaceOp(op, reduce_op.getResults());

    return success();
  }
};

// Convert tf.XlaVariadicSort to mhlo.Sort
class ConvertXlaVariadicSortOp
    : public OpRewritePattern<TF::XlaVariadicSortOp> {
 public:
  using OpRewritePattern::OpRewritePattern;

    Variadic<TF_Tensor>:$outputs
  );

  TF_DerivedOperandTypeListAttr T = TF_DerivedOperandTypeListAttr<0>;

  let hasVerifier = 1;
}

def TF_XlaVariadicSortOp : TF_Op<"XlaVariadicSort", [Pure]> {
  let summary = "Wraps the XLA Sort operator, documented at";

  let description = [{
https://www.tensorflow.org/performance/xla/operation_semantics#sort
.