let constructor = "TF::CreateHoistLoopInvariantPass()";
}

def LocalizeVarHandlesPass : Pass<"tf-localize-var-handles", "ModuleOp"> {
  let summary = "Creates VarHandleOps next to the operations that use them.";
  let description = [{
      Creates VarHandleOps right next to the operations that use them, one
      per operation.
      This is useful for transformations that only end up with a few small

      if (auto global = tf_saved_model::LookupBoundInputOfType<
              tf_saved_model::GlobalTensorOp>(func, i, syms)) {
        OpBuilder builder(func.getBody());
        auto dummy = builder.create<TF::VarHandleOp>(
            global.getLoc(), func.getArgument(i).getType(), "dummy", "dummy");
        func.getArgument(i).replaceAllUsesWith(dummy.getResult());
        argsToErase.set(i);
      }
    }

TFE_TensorHandle* TestVariable(TFE_Context* ctx, float value,
                               const tensorflow::string& device_name) {
  TF_Status* status = TF_NewStatus();
  // Create the variable handle.
  TFE_Op* op = TFE_NewOp(ctx, "VarHandleOp", status);
  if (TF_GetCode(status) != TF_OK) return nullptr;
  TFE_OpSetAttrType(op, "dtype", TF_FLOAT);
  TFE_OpSetAttrShape(op, "shape", {}, 0, status);
  TFE_OpSetAttrString(op, "container", "localhost", 0);

// aliasing output arguments.
std::unique_ptr<OperationPass<ModuleOp>> CreatePromoteResourcesToArgsPass(
    llvm::ArrayRef<std::string> functions = {});

// Creates a pass that promotes tf.VarHandleOp to resource arguments for all
// functions.
std::unique_ptr<OperationPass<ModuleOp>> CreatePromoteVarHandlesToArgsPass();

// Creates a pass that converts readonly reference variables to the

// Tests loop with two resource types, one of them being unique per iteration.
//
// Similar to above test but with one additional resource that is not unique per
// iteration (created by `tf.VarHandleOp`).

func.func @mixed_unique_resource_chain(%arg0: tensor<i32>, %arg1: tensor<f32>) {
  tf_executor.graph {

  if (auto pop_back = dyn_cast<TF::TensorListPopBackOp>(op)) {
    return InferShapeForTensorListPopBackOp(pop_back);
  }

  if (auto var_handle_op = dyn_cast<VarHandleOp>(op)) {
    return InferShapeForVarHandleOp(var_handle_op);
  }

  if (auto xla_reduce_window_op = dyn_cast<XlaReduceWindowOp>(op)) {
    return InferShapeForXlaReduceWindowOp(xla_reduce_window_op);
  }

  "tf_device.cluster"() ({
    // CHECK:      "tf.VarHandleOp"
    // CHECK-SAME: _xla_outside_compilation
    "tf.VarHandleOp"() {allowed_devices = [], container = "", device = "", shared_name = ""} : () -> tensor<!tf_type.resource<tensor<!tf_type.string>>>
    tf_device.return

  func.func @infer_var_handle_op_from_assigns() -> tensor<1xi8> {
    %cst = arith.constant dense<1> : tensor<1xi8>
    %0 = "tf.VarHandleOp"() {container = "", shared_name = "bar"} : () -> tensor<!tf_type.resource<tensor<*xi8>>>
    // CHECK: "tf.VarHandleOp"() <{container = "", shared_name = "bar"}> : () -> tensor<!tf_type.resource<tensor<1xi8>>>

  func.return %arg : tensor<!tf_type.resource<tensor<i32>>>
}

// CHECK-LABEL: func @side_effect_while_test
func.func @side_effect_while_test() -> (tensor<i32>) {
  %0 = "tf.VarHandleOp"() {device = "/device:CPU:0", container = "c", shared_name = "v"} : () -> tensor<!tf_type.resource<tensor<i32>>>
  // CHECK: [[while_res:%.]]:2 = tfrt.while {{%.*}} @"side_effect_while_body_add2/tfrt_body_1"

  return VerifyUnsortedSegmentReduction(*this);
}

//===----------------------------------------------------------------------===//
// VarHandleOp
//===----------------------------------------------------------------------===//

LogicalResult VarHandleOp::verify() {
  // VarHandleOp requires the resource handle supply a single subtype from
  // which to derive the dtype and shape attributes.