// CHECK: "tf.AssignVariableOp"(%[[TA_BUFFER]], %[[UPDATE]])
      // CHECK-NOT: tf.TensorArrayWriteV3
      %idx = "tf.Const"() {value = dense<4> : tensor<i32>} : () -> tensor<i32>
      %elem = "tf._SomeOp"() : () -> tensor<3xf32>
      %write = "tf.TensorArrayWriteV3"(%ta#0, %idx, %elem, %ta#1) : (tensor<!tf_type.resource>, tensor<i32>, tensor<3xf32>, tensor<f32>) -> tensor<f32>

  %y = "tf.MatMul"(%x, %cst) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?x100xf32>, tensor<100x512xf32>) -> (tensor<?x512xf32>)
  %flow_1_out = "tf.TensorArrayWriteV3"(%handle_1, %index, %y, %flow_1) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<2x!tf_type.resource<tensor<?x512xf32>>>, tensor<i32>, tensor<?x512xf32>, tensor<f32>) -> tensor<f32>

  // CHECK: tfrt_fallback_async.executeop key({{.*}}) cost({{.*}}) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.TensorArrayWriteV3"
  // CHECK: tfrt_fallback_async.executeop key({{.*}}) cost({{.*}}) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.TensorArrayGatherV3"

// CHECK-NEXT: [[flow_in_1:%.*]] = "tf_mlrt.tf_await"(%arg2) : (!mlrt.future) -> tensor<f32>
// CHECK-NEXT: [[flow_out_0:%.*]] = "tf.TensorArrayWriteV3"(%arg7, %arg5, [[result_0]], [[flow_in_0]])
// CHECK-NEXT: [[flow_out_1:%.*]] = "tf.TensorArrayWriteV3"(%arg8, %arg5, [[result_1]], [[flow_in_1]])
// CHECK-NEXT: "tf_mlrt.tf_promise"(%arg1, [[flow_out_0]]) : (!mlrt.promise, tensor<f32>) -> ()

 StackV2 314
 StridedSlice 12
 Sub 5
 Sum 2
 Switch 247
 TensorArrayGatherV3 8
 TensorArrayGradV3 17
 TensorArrayReadV3 9
 TensorArrayScatterV3 8
 TensorArraySizeV3 4
 TensorArrayV3 8
 TensorArrayWriteV3 9
 Unique 2
 VariableV2 164
 _Retval 5
cluster 0 size 440
 Abs 40
 AddN 1
 Any 41
 Cast 40
 ConcatV2 2
 Const 95
 ExpandDims 2
 IsInf 1
 IsNan 40
 L2Loss 40
 LogicalOr 1

    // safe.
    if (auto tensor_array_v3_op = dyn_cast<TF::TensorArrayV3Op>(tf_op)) {
      // The "flow" in TensorArrayV3 is always a scalar float tensor.
      // https://www.tensorflow.org/api_docs/python/tf/raw_ops/TensorArrayWriteV3
      Value flow = tensor_array_v3_op.getFlow();
      Type scalar_f32_tensor_type =
          RankedTensorType::get(/*shape=*/{}, rewriter.getF32Type());
      flow.setType(scalar_f32_tensor_type);
    }

  %y = "tf.MatMul"(%x, %cst) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<?x100xf32>, tensor<100x512xf32>) -> (tensor<?x512xf32>)
  %flow_1_out = "tf.TensorArrayWriteV3"(%handle_1, %index, %y, %flow_1) {device = "/job:localhost/replica:0/task:0/device:CPU:0"} : (tensor<2x!tf_type.resource<tensor<?x512xf32>>>, tensor<i32>, tensor<?x512xf32>, tensor<f32>) -> tensor<f32>

      "TensorArrayGatherV3",
      "TensorArrayGradV3",
      "TensorArrayReadV3",
      "TensorArrayScatterV3",
      "TensorArraySizeV3",
      "TensorArraySplitV3",
      "TensorArrayV3",
      "TensorArrayWriteV3",
      "TensorListConcatV2",
      "TensorListElementShape",
      "TensorListFromTensor",
      "TensorListGather",
      "TensorListGetItem",
      "TensorListLength",
      "TensorListPopBack",

    Res<TF_Float32Tensor, [{A scalar used to control gradient flow.}]>:$flow
  );
}

def TF_TensorArrayWriteV3Op : TF_Op<"TensorArrayWriteV3", []> {
  let summary = "Push an element onto the tensor_array.";

  let arguments = (ins
    Arg<TF_ResourceTensor, [{The handle to a TensorArray.}], [TF_TensorArrayRead, TF_TensorArrayWrite]>:$handle,