"tfl.assign_variable"(%0, %1) : (tensor<!tf_type.resource>, tensor<1x2x1x3xf32>) -> ()
  func.return %arg0 : tensor<1x2x1x3x!quant.uniform<i8:f32, 1.0>>

// CHECK-NEXT:  %[[vh:.*]] = "tfl.var_handle"() <{container = "", shared_name = ""}> : () -> tensor<*x!tf_type.resource<tensor<1x2x1x3x!quant.uniform<i8:f32, 1.000000e+00>>>>

// RUN: tf-tfrt-opt -tf-to-tfrt %s | FileCheck %s --dump-input=fail

// CHECK-LABEL: func @assign_variable
// CHECK-SAME: ([[in_chain:%.*]]: !tfrt.chain) -> !tfrt.chain
func.func @assign_variable() {
  // CHECK: [[ch1:%.*]], %results = tfrt_fallback_async.executeop.seq([[in_chain]]) key(0) cost({{.*}}) device("/device:CPU:0") "tf.VarHandleOp"

# RUN: tf-mlir-translate -graphdef-to-mlir -tf-enable-shape-inference-on-import=false %s -tf-graph-as-function -tf-control-output-arrays=assign_variable | tf-mlir-translate -mlir-tf-graph-to-hlo-text -tf-input-shapes=2:2 -tf-input-data-types=DT_FLOAT,DT_FLOAT -tf-xla-input-types=parameter,resource -tf-xla-emit-return-tuple | FileCheck %s

node {
  name: "arg0"
  op: "_Arg"
  attr {
    key: "T"
    value {
      type: DT_FLOAT
    }
  }
  attr {

  // CHECK: %[[ADD:.*]] = tfl.add %[[VAR_VAL]], %arg0 {fused_activation_function = "NONE"} : tensor<1x10xf32>
  // CHECK: "tfl.assign_variable"(%[[RESOURCE]], %[[ADD]]) : (tensor<!tf_type.resource<tensor<1x10xf32>>>, tensor<1x10xf32>) -> ()
  // CHECK: %[[RESULT:.*]] = "tfl.read_variable"(%[[RESOURCE]]) : (tensor<!tf_type.resource<tensor<1x10xf32>>>) -> tensor<1x10xf32>

    return success();
  }

  region.walk([&](Operation* op) {
    if (isa<TF::ReadVariableOp, func::ReturnOp, YieldOp>(op)) {
      return;
    }
    if (auto assign_variable = dyn_cast<TF::AssignVariableOp>(op)) {
      SetPotentiallyWritten(assign_variable.getResource());
      return;
    }
    if (auto call = dyn_cast<CallOpInterface>(op)) {
      if (auto func = dyn_cast<func::FuncOp>(call.resolveCallable())) {

  HASHTABLE = 136,
  HASHTABLE_FIND = 137,
  HASHTABLE_IMPORT = 138,
  HASHTABLE_SIZE = 139,
  REDUCE_ALL = 140,
  CONV_3D_TRANSPOSE = 141,
  VAR_HANDLE = 142,
  READ_VARIABLE = 143,
  ASSIGN_VARIABLE = 144,
}
// LINT.ThenChange(nnapi_linter/linter.proto)

// Options for the builtin operators.
union BuiltinOptions {
  Conv2DOptions,
  DepthwiseConv2DOptions,
  ConcatEmbeddingsOptions,

  HASHTABLE = 136,
  HASHTABLE_FIND = 137,
  HASHTABLE_IMPORT = 138,
  HASHTABLE_SIZE = 139,
  REDUCE_ALL = 140,
  CONV_3D_TRANSPOSE = 141,
  VAR_HANDLE = 142,
  READ_VARIABLE = 143,
  ASSIGN_VARIABLE = 144,
  BROADCAST_ARGS = 145,
  RANDOM_STANDARD_NORMAL = 146,
  BUCKETIZE = 147,
  RANDOM_UNIFORM = 148,
  MULTINOMIAL = 149,
  GELU = 150,
  DYNAMIC_UPDATE_SLICE = 151,
  RELU_0_TO_1 = 152,

    DefaultValuedStrAttr<StrAttr, "">:$shared_name
  );

  let results = (outs TFL_ResourceTensor:$resource_handle);

  let hasOptions = 1;
}

def TFL_AssignVariableOp : TFL_Op<"assign_variable", []> {
  let summary = "Assigns a new value to a variable.";

  let description = [{
Any ReadVariableOp with a control dependency on this op is guaranteed to return

    "HASHTABLE",
    "HASHTABLE_FIND",
    "HASHTABLE_IMPORT",
    "HASHTABLE_SIZE",
    "REDUCE_ALL",
    "CONV_3D_TRANSPOSE",
    "VAR_HANDLE",
    "READ_VARIABLE",
    "ASSIGN_VARIABLE",
    "BROADCAST_ARGS",
    "RANDOM_STANDARD_NORMAL",
    "BUCKETIZE",
    "RANDOM_UNIFORM",
    "MULTINOMIAL",
    "GELU",
    "DYNAMIC_UPDATE_SLICE",
    "RELU_0_TO_1",