"tf.AssignVariableOp"(%arg3, %13#10) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<128x1024xf32>>>, tensor<128x1024xf32>) -> ()
    "tf.AssignVariableOp"(%arg5, %13#8) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<1024x1xf32>>>, tensor<1024x1xf32>) -> ()

// CHECK: "tf.AssignVariableOp"(%[[VAR_HANDLE_0]], %[[RESTORE]]#0) <{validate_shape = false}> : (tensor<!tf_type.resource<tensor<2xf32>>>, tensor<2xf32>) -> ()
// CHECK: "tf.AssignVariableOp"(%[[VAR_HANDLE_1]], %[[RESTORE]]#1) <{validate_shape = false}> : (tensor<!tf_type.resource<tensor<4xi32>>>, tensor<4xi32>) -> ()

// Test that the locs are properly set to their shared_names.

// conditions are met:
//   1. `validate_shape` is set to `true`
//   2. The variable is initialized
//   3. The shape of the value tensor doesn't match the shape of the variable
//      tensor.
TF_CAPI_EXPORT extern void TF_AssignVariable(
    TF_OpKernelContext* ctx, int input_index, int value_index,
    bool validate_shape,
    void (*copyFunc)(TF_OpKernelContext* ctx, TF_Tensor* source,

  // CHECK: "tf.AssignVariableOp"([[VAR_HANDLE]], [[NEW_VAR]]) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<*xf32>>>, tensor<*xf32>) -> ()
  // CHECK: "tf.AssignVariableOp"([[M_HANDLE]], [[NEW_M]]) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<*xf32>>>, tensor<*xf32>) -> ()
  // CHECK: "tf.AssignVariableOp"([[V_HANDLE]], [[NEW_V]]) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<*xf32>>>, tensor<*xf32>) -> ()

      return nullptr;
    }
  }
  return cc_ctx->input_ref_mutex(input);
}

void TF_AssignVariable(TF_OpKernelContext* ctx, int input_index,
                       int value_index, bool validate_shape,
                       void (*copyFunc)(TF_OpKernelContext* ctx,
                                        TF_Tensor* source, TF_Tensor* dest),
                       TF_Status* status) {

    "tf.AssignVariableOp"(%arg0, %3) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<1xf32>>>, tensor<1xf32>) -> ()
    tf_device.return
  }
  func.return
}

func.func private @write_chain_func(%arg0: tensor<1xf32>) -> (tensor<1xf32>) {

    %0 = "tf.VarHandleOp"() <{container = "", shared_name = "Variable"}> : () -> tensor<!tf_type.resource<tensor<!tf_type.string>>>
    "tf.AssignVariableOp"(%0, %cst) <{validate_shape = false}> : (tensor<!tf_type.resource<tensor<!tf_type.string>>>, tensor<!tf_type.string>) -> ()
    %2 = "tf.ReadVariableOp"(%0) : (tensor<!tf_type.resource<tensor<!tf_type.string>>>) -> tensor<*x!tf_type.string>

    Arg<TF_Tensor, [{Should be from a `Variable` node. May be uninitialized.}]>:$ref,
    Arg<TF_Tensor, [{The value to be assigned to the variable.}]>:$value,

    DefaultValuedOptionalAttr<BoolAttr, "true">:$validate_shape,
    DefaultValuedOptionalAttr<BoolAttr, "true">:$use_locking
  );

  let results = (outs
    Res<TF_Tensor, [{= Same as "ref".  Returned as a convenience for operations that want

      }
    }

    auto set_shape_from_list_attr = [&](const AttrValue* attr) {
      auto& list = attr->list();
      // This follows the same approach as in ValidateShape, but only flags
      // warning in case where there are mismatch in number of shapes and
      // outputs and in which case it just returns without attempting to refine.
      if (list.shape_size() != node->num_outputs()) {