// shown above jit/partially_decluster_pass extracts the actual_size(...)
// computation to outside the XLA cluster, causing the cluster to be versioned
// only on the actual size of the XlaDynamicSlice.  This avoids recompilation
// due to superficial changes that don't affect tensor shapes.
//
// Future Work TODO(b/111210515)
// -----------------------------
//

                          XlaCompileOnDemandOp);                               \
  REGISTER_KERNEL_BUILDER(                                                     \
      Name("XlaDynamicSlice").HostMemory("size_indices").Device(DEVICE),       \
      XlaCompileOnDemandOp);                                                   \
  REGISTER_KERNEL_BUILDER(Name("XlaDynamicUpdateSlice").Device(DEVICE),        \

          [(HasValidDotDims $dimension_numbers), (HasValidPrecisionConfig $precision_config)]>;

//===----------------------------------------------------------------------===//
// XlaDynamicSlice op.
//===----------------------------------------------------------------------===//

def : Pat<(TF_XlaDynamicSliceOp:$op MHLO_Tensor:$input, MHLO_Tensor:$starting_indices,

                                         "XlaDequantize",
                                         "XlaDot",
                                         "XlaDotV2",
                                         "XlaDynamicSlice",
                                         "XlaDynamicUpdateSlice",
                                         "XlaEinsum",
                                         "XlaGather",

  %0 = "tf.XlaDynamicSlice"(%arg0, %arg1, %sizes) : (tensor<3x4xi32>, tensor<2xi64>, tensor<2xi64>) -> tensor<1x4xi32>
  func.return %0 : tensor<1x4xi32>
}

// -----

// CHECK-LABEL: xla_dynamic_slice_mhlo_sizes
func.func @xla_dynamic_slice_mhlo_sizes(%arg0: tensor<1x1024x4xf32>, %arg1: tensor<3xi32>) -> tensor<1x512x4xf32> {
  // CHECK-NOT: "tf.XlaDynamicSlice"

      "XlaConcatND",
      "XlaConv",
      "XlaConvV2",
      "XlaCustomCall",
      "XlaCustomCallV2",
      "XlaDequantize",
      "XlaDot",
      "XlaDotV2",
      "XlaDynamicSlice",
      "XlaDynamicUpdateSlice",
      "XlaEinsum",
      "XlaGather",
      "XlaIf",
      "XlaKeyValueSort",
      "XlaOptimizationBarrier",
      "XlaPad",
      "XlaRecv",

  TF_DerivedOperandTypeAttr RhsT = TF_DerivedOperandTypeAttr<1>;
  TF_DerivedResultTypeAttr preferred_element_type = TF_DerivedResultTypeAttr<0>;
}

def TF_XlaDynamicSliceOp : TF_Op<"XlaDynamicSlice", [Pure, TF_NoConstantFold]> {
  let summary = "Wraps the XLA DynamicSlice operator, documented at";

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