// 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);