using tensorflow::FunctionDef;
using tensorflow::Node;
using tensorflow::NodeBuilder;
using tensorflow::Status;
using tensorflow::errors::InvalidArgument;

namespace {
typedef std::unique_ptr<TF_Function, decltype(&TF_DeleteFunction)>
    UniqueFuncPtr;
}

// struct TF_Operation { tensorflow::Node node; };
static TF_Operation* ToTF_Operation(Node* node) {
  return static_cast<TF_Operation*>(static_cast<void*>(node));
}

  auto* tmp_var = new TmpVar;
  OP_REQUIRES(context, tmp_var,
              absl::ResourceExhaustedError("Could not allocate TmpVar."));
  tmp_var->name = unique_name;

  Status s;
  std::unique_ptr<TF_Tensor, decltype(&TF_DeleteTensor)> tmp_var_tf(
      tensorflow::TF_TensorFromTensor(tmp_var->val, &s), TF_DeleteTensor);
  OP_REQUIRES_OK(context, s);
  allocFunc(ctx, tmp_var_tf.get(), dtype, dims, num_dims, tf_status);

std::optional<StringRef> GetShardingStringFromVariant(
    const OpShardingVariant& sharding_or_op) {
  return std::visit(
      [](auto&& sharding_or_op) -> std::optional<StringRef> {
        using T = std::decay_t<decltype(sharding_or_op)>;
        if constexpr (std::is_same_v<T, StringRef>) {
          return sharding_or_op;
        } else {
          return GetXlaShardingFromOperator(sharding_or_op);
        }
      },

    if (methods_.summarize == nullptr) {
      return tensorflow::CustomDeviceTensorHandle::SummarizeValue(summary);
    }
    TF_Status c_status;
    std::unique_ptr<TF_Buffer, decltype(&TF_DeleteBuffer)> summary_buffer(
        methods_.summarize(data_, &c_status), TF_DeleteBuffer);
    if (!c_status.status.ok()) {
      return c_status.status;
    }