func.func @tfInplaceUpdate(%arg0: tensor<2x1x2xf32>) -> tensor<2x1x2xf32> {
  %1 = arith.constant dense<1> : tensor<1xi32>
  %2 = arith.constant dense<2.0> : tensor<1x1x2xf32>
  %3 = "tf.InplaceUpdate"(%arg0, %1, %2) {device = ""}
    : (tensor<2x1x2xf32>, tensor<1xi32>, tensor<1x1x2xf32>) -> tensor<2x1x2xf32>
  func.return %3 : tensor<2x1x2xf32>
}

}

// CHECK-LABEL: @tfInplaceUpdate

module {
func.func @tfInplaceUpdate(%arg0: tensor<2x1x2xf32>) -> tensor<2x1x2xf32> {
  %1 = arith.constant dense<1> : tensor<1xi32>
  %2 = arith.constant dense<2.0> : tensor<1x1x2xf32>
  %3 = "tf.InplaceUpdate"(%arg0, %1, %2) {device = ""}
    : (tensor<2x1x2xf32>, tensor<1xi32>, tensor<1x1x2xf32>) -> tensor<2x1x2xf32>
  func.return %3 : tensor<2x1x2xf32>
}
}

//CHECK: module attributes

module {
func.func @tfInplaceUpdate(%arg0: tensor<2x1x2xf32>) -> tensor<2x1x2xf32> {
  %1 = arith.constant dense<1> : tensor<1xi32>
  %2 = arith.constant dense<2.0> : tensor<1x1x2xf32>
  %3 = "tf.InplaceUpdate"(%arg0, %1, %2) {device = ""}
    : (tensor<2x1x2xf32>, tensor<1xi32>, tensor<1x1x2xf32>) -> tensor<2x1x2xf32>
  func.return %3 : tensor<2x1x2xf32>
}
}

//CHECK: module {

    "FakeQuantWithMinMaxVarsPerChannelGradient",
    "Fill",
    "Fingerprint",
    "Gather",
    "GuaranteeConst",
    "Identity",
    "IdentityN",
    "InplaceAdd",
    "InplaceSub",
    "InplaceUpdate",
    "InvertPermutation",
    "ListDiff",
    "MatrixBandPart",
    "MatrixDiag",
    "MatrixDiagPart",
    "MatrixDiagPartV2",
    "MatrixDiagPartV3",
    "MatrixDiagV2",
    "MatrixDiagV3",

absl::StatusOr<bool> IsSmallConstant(Node* n) {
  if (!n->IsConstant()) {
    return false;
  }

  return IsConstantSmall(n);
}

bool IsInPlaceOp(absl::string_view op_name) {
  return op_name == "InplaceUpdate" || op_name == "InplaceAdd" ||
         op_name == "InplaceSub";
}
}  // namespace

Status CloneConstantsForBetterClusteringPassImpl::CloneSmallConstantInputs(

  // CHECK-DAG: [[RESHAPE1:%.+]] = mhlo.reshape [[SLICE1]]
  // CHECK-DAG: [[UPDATE:%.+]] = mhlo.dynamic_update_slice %arg0, [[SLICE2]], [[RESHAPE1]], [[CST]]
  %0 = "tf.InplaceUpdate"(%arg0, %arg2, %arg1) : (tensor<8x4xf32>, tensor<1xi32>, tensor<1x4xf32>) -> tensor<8x4xf32>

  // CHECK: return [[UPDATE]]
  func.return %0 : tensor<8x4xf32>
}

// -----

    custom_call->setAttr(kShardingAttr, op.get_XlaShardingAttr());
    rewriter.replaceOp(op, custom_call.getResult(0));

    return success();
  }
};

// Converts a TF InplaceUpdate op to DynamicUpdateSlice HLO.
class ConvertInplaceUpdateOp : public OpRewritePattern<TF::InplaceUpdateOp> {
 public:
  using OpRewritePattern::OpRewritePattern;

  );

  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
}

def TF_InplaceUpdateOp : TF_Op<"InplaceUpdate", [Pure]> {
  let summary = "Updates specified rows 'i' with values 'v'.";

  let description = [{
Computes `x[i, :] = v; return x`.