"defs.h",
    ],
    visibility = [":friends"],
)

# Internal targets below this point.

cc_library(
    name = "variable_info",
    srcs = ["variable_info.cc"],
    hdrs = ["variable_info.h"],
    visibility = [
        ":internal",
        # We reuse VariableInfo in TFRT's implementation of TpuExecuteOp.
        "//learning/brain/tfrt/tf_tpu:__pkg__",
        "//learning/brain/tfrt/tpu_plugin:__pkg__",

          *ptr = new Var(DT_INVALID);
          return absl::OkStatus();
        }));
    VariableInfo& variable_info = result->emplace_back(
        var_idx, handle.name(), variable, handle.definition_stack_trace());
    if (variables_updated != nullptr &&
        variables_updated->find(var_idx) == variables_updated->end()) {
      variable_info.set_read_only();
    }
  }
  return absl::OkStatus();
}

  absl::StatusOr<std::vector<VariableInfo>> variable_infos =
      GatherVariableInfo(ctx, *result, 0);
  TF_RETURN_IF_ERROR(variable_infos.status());
  TF_RETURN_IF_ERROR(LockVariables(absl::MakeSpan(*variable_infos)));
  TF_RETURN_IF_ERROR(launch_context.PopulateOutputs(
      ctx, result, execution_output.ConsumeResult(),
      /*missing_ctx_input_prefix=*/0, absl::MakeSpan(*variable_infos),
      input_output_alias, snapshot_ptrs));

  // operate on TF_VariableInfo (such as allocate temp tensor for the `var` held
  // by the underlying tensorflow::VariableInfo.
  TF_VariableInfo(int index, const std::string& name, tensorflow::Var* var) {
    var_info = tensorflow::VariableInfo{index, name, var};
  }

  tensorflow::VariableInfo var_info{0, "", nullptr};
};

TF_VariableInfo* TF_CreateVariableInfoFromContext(TF_OpKernelContext* ctx,

  Status PopulateOutputs(
      OpKernelContext* ctx,
      const XlaCompiler::CompilationResult* compilation_result,
      xla::ScopedShapedBuffer output, int missing_ctx_input_prefix,
      absl::Span<VariableInfo> variable_infos,
      const xla::HloInputOutputAliasConfig& input_output_alias,
      const std::map<int, const Tensor*>& resource_vars);

 private:
  xla::LocalClient* client_;

                                                  inputs, resource_arg_indices,
                                                  &variable_infos));
    TF_RETURN_IF_ERROR(LockVariables(absl::MakeSpan(variable_infos)));
  }

  absl::StatusOr<std::vector<XlaCompiler::Argument>> args;

  if (compiler_arg_source == CompilerArgSource::TENSOR_SPEC) {

  std::pair<std::vector<XlaCompiler::Argument>, ResourceVarsSnapshot> result;

  std::vector<VariableInfo> variable_infos;
  TF_RETURN_IF_ERROR(
      GetVariableInfosFromInputs(ctx->resource_manager(), ctx->device(), inputs,
                                 variable_indices, &variable_infos));
  TF_RETURN_IF_ERROR(LockVariables(absl::MakeSpan(variable_infos)));

  TF_RETURN_IF_ERROR(SnapshotResourceVariables(ctx, variable_indices,

      ctx->set_output(i, output_tensor);
      ++output_num;
    }
  }

  // input_index -> index into variable_infos.
  absl::flat_hash_map<int, int> variable_info_lookup;
  for (int i = 0; i < variable_infos.size(); i++) {
    variable_info_lookup.emplace(variable_infos[i].index(), i);
  }

  // Apply variable updates, if any.
  for (int i = 0, end = compilation_result->resource_updates.size(); i < end;

#include "tensorflow/compiler/jit/device_compiler.h"
#include "tensorflow/compiler/jit/flags.h"
#include "tensorflow/compiler/jit/pjrt_device_compiler_client.h"
#include "tensorflow/compiler/jit/variable_info.h"
#include "tensorflow/compiler/jit/variable_info_util.h"
#include "tensorflow/compiler/tf2xla/xla_op_registry.h"
#include "xla/pjrt/pjrt_client.h"
#include "xla/pjrt/tfrt_cpu_pjrt_client.h"