}

func.func @dynamic_update_slice(%arg0: tensor<4x5xi32>, %arg1: tensor<1x5xi32>, %arg2: tensor<2xi32>) -> tensor<4x5xi32> {
  %0 = "tf.XlaDynamicUpdateSlice"(%arg0, %arg1, %arg2) : (tensor<4x5xi32>, tensor<1x5xi32>, tensor<2xi32>) -> tensor<4x5xi32>
  func.return %0 : tensor<4x5xi32>

// CHECK-LABEL:dynamic_update_slice

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

// CHECK:         %0 = "tf.Pack"(%arg2, %arg3, %arg4) <{axis = 0 : i64}> : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<3xi32>
// CHECK:         %1 = "tf.XlaDynamicUpdateSlice"(%arg0, %arg1, %0) : (tensor<28x1x100xf32>, tensor<1x1x100xf32>, tensor<3xi32>) -> tensor<28x1x100xf32>
// CHECK:         return %1 : tensor<28x1x100xf32>

  // CHECK: [[DUS:%.+]] = mhlo.dynamic_update_slice %arg0, %arg1, [[RESHAPE0]], [[RESHAPE1]] : (tensor<4x16xf32>, tensor<2x4xf32>, tensor<i32>, tensor<i32>) -> tensor<4x16xf32>
  // CHECK: return [[DUS]]
  %0 = "tf.XlaDynamicUpdateSlice"(%arg0, %arg1, %arg2) : (tensor<4x16xf32>, tensor<2x4xf32>, tensor<2xi32>) -> tensor<4x16xf32>
  func.return %0 : tensor<4x16xf32>
}

// -----

// CHECK-LABEL: xla_dynamic_update_slice2

          op.getLoc(), op.getType(), input, std::get<1>(pair), input_indices);
    }

    rewriter.replaceOp(op, input);
    return success();
  }
};

// Converts a TF XlaDynamicUpdateSlice op to DynamicUpdateSlice HLO.
class ConvertXlaDynamicUpdateSliceOp
    : public OpRewritePattern<TF::XlaDynamicUpdateSliceOp> {
 public:
  using OpRewritePattern::OpRewritePattern;

}

def TF_XlaDynamicUpdateSliceOp : TF_Op<"XlaDynamicUpdateSlice", [Pure, TF_NoConstantFold]> {
  let summary = "Wraps the XLA DynamicUpdateSlice operator, documented at";

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

XlaDynamicUpdateSlice generates a result which is the value of the `input`