%index = "tf.Const"() {value = dense<1> : tensor<i32>} : () -> tensor<i32>
  // CHECK: %[[VAR:.*]] = "tf.MlirLocalVarOp"() : () -> tensor<!tf_type.resource<tensor<5x3xf32>>>
  // CHECK: %[[GVAR:.*]] = "tf.MlirLocalVarOp"() : () -> tensor<!tf_type.resource<tensor<5x3xf32>>>

==============================================================================*/

#include "tensorflow/core/framework/common_shape_fns.h"
#include "tensorflow/core/framework/op.h"

namespace tensorflow {

REGISTER_OP("MlirLocalVarOp")
    .Output("resource: resource")
    .SetShapeFn(shape_inference::UnknownShape)
    .Doc(R"(Creates a handle to an in-scope variable.
Used by internal passes for temporary representation of local state, which will

  // CHECK-NEXT: %[[BROADCAST:.*]] = "tf.BroadcastTo"(%[[CAST_ZERO]], %[[CONST10]]) : (tensor<f32>, tensor<1xi32>) -> tensor<10xf32>
  // CHECK-NEXT: %[[BUFFER:.*]] = "tf.MlirLocalVarOp"() : () -> tensor<!tf_type.resource<tensor<10xf32>>>
  // CHECK-NEXT: %[[SIZE:.*]] = "tf.MlirLocalVarOp"() : () -> tensor<!tf_type.resource<tensor<1xi32>>>
  // CHECK-NEXT: %[[ZERO:.*]] = "tf.Const"() <{value = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32>

    var_handle_op->erase();
  } else {
    std::get<TF::MlirLocalVarOp>(op).erase();
  }
}

std::optional<LocalVarOp> IsLocalVarOp(Operation &op) {
  if (TF::MlirLocalVarOp mlir_local_var_op =
          llvm::dyn_cast<TF::MlirLocalVarOp>(&op)) {
    return std::make_optional(LocalVarOp(mlir_local_var_op));
  }
  if (TF::VarHandleOp var_handle_op = llvm::dyn_cast<TF::VarHandleOp>(&op)) {

  func.return
}

// -----

// Tests removal of dead local variables.

// CHECK-LABEL: func @main
func.func @main(%arg0: tensor<2xf32>) {
  // CHECK-NOT: tf.MlirLocalVarOp
  // CHECK-NOT: tf.AssignVariableOp
  %0 = "tf.MlirLocalVarOp"() : () -> tensor<!tf_type.resource<tensor<2xf32>>>
  "tf.AssignVariableOp"(%0, %arg0) : (tensor<!tf_type.resource<tensor<2xf32>>>, tensor<2xf32>) -> ()
  func.return
}

// -----

              ArrayRef<TensorType>{mlir::cast<TensorType>(buffer.getType())},
              stack.getContext()));
  auto local_var = builder.create<TF::MlirLocalVarOp>(
      stack.getLoc(), ArrayRef<Type>{var_type}, ArrayRef<Value>{});
  auto local_size_var = builder.create<TF::MlirLocalVarOp>(
      stack.getLoc(), ArrayRef<Type>{size_var_type}, ArrayRef<Value>{});
  // Zero-initialize the local vars.
  cutil::WriteLocalVariable(local_size_var,

  auto var_type = RankedTensorType::get(
      {}, TF::ResourceType::get(
              ArrayRef<TensorType>{buffer.getType().cast<TensorType>()},
              ta.getContext()));
  auto local_var = builder.create<TF::MlirLocalVarOp>(
      ta.getLoc(), ArrayRef<Type>{var_type}, ArrayRef<Value>{});
  cutil::WriteLocalVariable(local_var, buffer, builder, ta.getLoc());
  ta.getHandle().replaceAllUsesWith(local_var);

  );

  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
  TF_DerivedOperandTypeAttr Tpaddings = TF_DerivedOperandTypeAttr<1>;
}

def TF_MlirLocalVarOp : TF_Op<"MlirLocalVarOp", []> {
  let summary = "Creates a handle to an in-scope variable.";

  let description = [{
Used by internal passes for temporary representation of local state, which will
be eventually removed.