const char* attr_name, const void* proto,
                                size_t proto_len, TF_Status* status) {
  // shape.ParseFromArray takes an int as length, this function takes size_t,
  // make sure there is no information loss.
  if (proto_len > std::numeric_limits<int>::max()) {
    status->status = InvalidArgument(
        "proto_len (", proto_len,

        "GradientDef is not supported in reading Dataset related functions: ",
        text_proto);
    return {};
  }
  std::vector<UniqueFuncPtr> ret;
  for (const FunctionDef& fdef : fdef_lib.function()) {
    // Make a copy so that we can mutate it.
    FunctionDef fdef_to_load = fdef;
    if (mutate_proto_func) {
      (*mutate_proto_func)(&fdef_to_load);
    }
    VLOG(1) << "Adding func to graph: " << fdef_to_load.DebugString();

    int64_t target_id = target_tensor_ids[i];
    if (sources_set.find(target_id) != sources_set.end()) {
      auto tensor = targets[i];
      sources_that_are_targets.insert(
          std::make_pair(target_id, TapeTensor(tensor)));
    }
  }

  TF_RETURN_IF_ERROR(GradientTape::ComputeGradient(
      vspace, target_tensor_ids, source_tensor_ids, sources_that_are_targets,

  CHECK_OK(buffer.status());

  return new PJRT_Buffer{std::move(*buffer), c_api_client->pjrt_c_client()};
}

TEST(TensorPjRtBufferUtilTest, GetPjRtCBufferFromTensorNoBuffer) {
  auto allocator = std::make_unique<AsyncValueAllocator>();
  tensorflow::Tensor tensor(allocator.get(), DT_FLOAT, {1});

  EXPECT_THAT(
      GetPjRtCBufferFromTensor(&tensor),
      StatusIs(error::INTERNAL, HasSubstr(absl::StrCat(

  TFE_TensorHandle* h0 = TestVariable(ctx, 1.0, task1_name);
  TFE_TensorHandle* h1 = TestVariable(ctx, 2.0, task2_name);
  TFE_TensorHandle* h2 = TestVariable(ctx, 3.0, task0_name);

  // Add a sync point to make sure that variables have been initialized
  // before the function execution starts.
  TFE_ContextAsyncWait(ctx, status);
  EXPECT_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);

  int num_retvals = 1;
  std::vector<AbstractTensorHandle*> outputs(1);
  GradientRegistry registry;
  s = RegisterGradients(&registry);
  ASSERT_EQ(errors::OK, s.code()) << s.message();
  auto tape = std::make_unique<Tape>(/*persistent=*/false);
  s = Execute(check_numerics_op.get(), ctx.get(), absl::MakeSpan(outputs),
              &num_retvals, &forward_op, tape.get(), registry);
  ASSERT_EQ(errors::OK, s.code()) << s.message();

ParallelDevice::ParallelDevice(const std::vector<std::string>& devices,
                               bool is_async, int in_flight_nodes_limit)
    : underlying_devices_(devices),
      default_cancellation_manager_(absl::make_unique<CancellationManager>()) {
  device_threads_.reserve(devices.size());
  for (int device_index = 0; device_index < devices.size(); ++device_index) {
    device_threads_.emplace_back(new DeviceThread(

}

// Converts `ninputs` and `inputs` into `inputs_tensors` and `input_nodes` and
// does various checks while doing so. `input_nodes` will contain the same
// information as input_tensors just in a different structure to make
// following processing easier. TODO(iga): Simplify this nested structure.
Status ProcessInputs(
    const TF_Graph* fn_body, const char* fn_name, int ninputs,
    const TF_Output* inputs, std::vector<OutputTensor>* input_tensors,

}

absl::Status SetPjRtCBufferToTensor(PJRT_Buffer* c_buffer,
                                    xla::PjRtCApiClient* c_api_client,
                                    Tensor* tensor) {
  auto buffer = std::make_unique<xla::PjRtCApiBuffer>(c_api_client, c_buffer);
  tensorflow::AsyncValueTensor* av_tensor =
      tensorflow::AsyncValueTensor::FromTensor(tensor);
  if (av_tensor == nullptr) {
    TF_ASSIGN_OR_RETURN(

      /*run_eager_op_as_function=*/opts->run_eager_op_as_function,
      /*jit_compile_rewrite=*/opts->jit_compile_rewrite);
#if !defined(IS_MOBILE_PLATFORM)
  eager_context->SetDistributedManager(
      std::make_unique<tensorflow::EagerContextDistributedManager>(
          eager_context));
#endif  // !IS_MOBILE_PLATFORM
  return tensorflow::wrap(eager_context);
}

void TFE_DeleteContext(TFE_Context* ctx) {