# RUN: tf-mlir-translate -graphdef-to-mlir -tf-upgrade-legacy %s -tf-output-arrays=hash_table_node,variable_node,variable_v2_node -o - | FileCheck %s

node: {
  name: "hash_table_node"
  op: "HashTableV2"
  attr: {
    key: "key_dtype"
    value: {
      type: DT_INT32
    }
  }
  attr: {
    key: "shared_name"
    value: {
      s: ""
    }
  }
  attr: {
    key: "value_dtype"
    value: {
      type: DT_FLOAT
    }

    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "shape"
    value {
      shape {
        dim {
          size: 1
        }
      }
    }
  }
}
node {
  name: "variable_node"
  op: "Const"
  input: "^input"
  attr {
    key: "dtype"
    value {
      type: DT_FLOAT
    }
  }
  attr {
    key: "value"
    value {
      tensor {
        dtype: DT_FLOAT

// Converts a Variable NodeDef into a Constant NodeDef.
void ConvertVariableToConstant(const NodeDef& variable_node,
                               const Tensor& variable_value,
                               NodeDef* const_node) {
  const_node->set_name(variable_node.name());
  const_node->set_op("Const");
  (*const_node->mutable_attr())["dtype"] = variable_node.attr().at("dtype");
  variable_value.AsProtoTensorContent(

}

// VariableNode holds a list of variable names, possibly with chained field
// accesses. The dollar sign is part of the (first) name.
type VariableNode struct {
	NodeType
	Pos
	tr    *Tree
	Ident []string // Variable name and fields in lexical order.
}

func (t *Tree) newVariable(pos Pos, ident string) *VariableNode {

// Create a variable handle with name `variable_name` on a device with name
// `device_name`.
TFE_TensorHandle* CreateVarHandle(TFE_Context* ctx,
                                  const tensorflow::string& device_name,
                                  const tensorflow::string& variable_name);

// Create a variable with value `value` and name `variable_name` on a device
// with name `device_name`.

'''
 Repository rule to set some environment variables.
 Can be set via build parameter "--repo_env=<VARIABLE_NAME>=<value>"
 e.g "--repo_env=REQUIREMENTS_FILE_NAME=requirements.in"

 List of variables:
 REQUIREMENTS_FILE_NAME
'''

def _updater_config_repository_impl(repository_ctx):
    repository_ctx.file("BUILD", "")
    requirements_file_name = repository_ctx.os.environ.get("REQUIREMENTS_FILE_NAME", "requirements.in")
    repository_ctx.file(

    StringAttr device_attr =
        variable_v2_op->getAttrOfType<StringAttr>("device");
    if (!device_attr) device_attr = builder.getStringAttr("");
    StringRef variable_name =
        GetNodeNameFromClassAttrOrSharedNameAttr(variable_v2_op);
    if (variable_name.empty()) {
      return signalPassFailure();
    }
    VarHandleOp var_handle_op = builder.create<VarHandleOp>(
        variable_v2_op.getLoc(),

  for (auto func_op : module.getOps<func::FuncOp>()) {
    for (auto var_handle_op : func_op.getOps<TF::VarHandleOp>()) {
      auto variable_name = GetVariableName(var_handle_op);
      if (variable_names.count(variable_name)) continue;
      var_ops.emplace_back(var_handle_op);
      variable_names.insert(variable_name);
    }
  }

  // Get resources from Session.
  auto resource_tensors_or = GetResourcesFromSession(var_ops, session);

  TFE_OpSetAttrShape(op, "shape", {}, 0, status);
  TFE_OpSetAttrString(op, "container", "localhost", 0);
  TFE_OpSetAttrString(op, "shared_name", variable_name.data(),
                      variable_name.size());
  if (!device_name.empty()) {
    TFE_OpSetDevice(op, device_name.c_str(), status);
  }
  if (TF_GetCode(status) != TF_OK) return nullptr;
  TFE_TensorHandle* var_handle = nullptr;
  int num_retvals = 1;

  // Read the variable using `ctx_1`. This read should succeed.
  // 1. Create a variable on `remote_device`, using `ctx_0`.
  TFE_TensorHandle* handle_0 =
      CreateVariable(ctx_0, 1.2, remote_device, /*variable_name=*/"var2");

  // 2. Wait for `var2` to be created and initialized on the worker.
  TF_Status* status = TF_NewStatus();
  TFE_ContextAsyncWait(ctx_0, status);