(*var_to_data_type_map)[key] = DataType(entry.dtype());
  }
  // The returned pointers are owned by the caller.
  return std::make_pair(std::move(var_to_shape_map),
                        std::move(var_to_data_type_map));
}

}  // namespace checkpoint

                               file.get(), plugin_status.get());

  if (TF_GetCode(plugin_status.get()) == TF_OK)
    *result = std::make_unique<ModularRandomAccessFile>(
        translated_name, std::move(file), random_access_file_ops_.get());

  return StatusFromTF_Status(plugin_status.get());
}

Status ModularFileSystem::NewWritableFile(
    const std::string& fname, TransactionToken* token,

      : filesystem_(std::move(filesystem)),
        ops_(std::move(filesystem_ops)),
        random_access_file_ops_(std::move(random_access_file_ops)),
        writable_file_ops_(std::move(writable_file_ops)),
        read_only_memory_region_ops_(std::move(read_only_memory_region_ops)),
        plugin_memory_allocate_(std::move(plugin_memory_allocate)),
        plugin_memory_free_(std::move(plugin_memory_free)) {}

  if (av_tensor == nullptr) {
    TF_ASSIGN_OR_RETURN(
        *tensor, MakeTensorFromPjRtBuffer(tensor->dtype(), tensor->shape(),
                                          std::move(buffer)));
  } else {
    av_tensor->SetBuffer(std::move(buffer));
  }
  return absl::OkStatus();
}

absl::StatusOr<xla::PjRtCApiClient*> GetPjRtCApiClient(
    const DeviceType& device_type) {

set -euxo pipefail

cd "$TFCI_OUTPUT_DIR"

# Move extra wheel files somewhere out of the way. This script
# expects just one wheel file to exist.
if [[ "$(ls *.whl | wc -l | tr -d ' ')" != "1" ]]; then
  echo "More than one wheel file is present: moving the oldest to"
  echo "$TFCI_OUTPUT_DIR/extra_wheels."
  # List all .whl files by their modification time (ls -t) and move anything

          {std::move(bucket), std::move(object), &gcs_file->gcs_client,
           TempFile(temp_file_name, std::ios::binary | std::ios::trunc), false,
           static_cast<int64_t>(metadata->size())});
    } else if (TF_GetCode(status) == TF_NOT_FOUND) {
      file->plugin_file = new tf_writable_file::GCSFile(
          {std::move(bucket), std::move(object), &gcs_file->gcs_client,

touch "$REQUIREMENTS_LOCK_FILE"

### Update the requirements_lock file
bazel run --experimental_convenience_symlinks=ignore --repo_env=REQUIREMENTS_FILE_NAME=requirements_wheel_test.in //:requirements_${PYTHON_VERSION}.update

# Move the updated file to the appropriate directory
mv "$REQUIREMENTS_LOCK_FILE" ../wheel_test/

    }
  }

  TF_DeleteTensor(analytical_tensor);
  delete[] danalytical;
}

Model BuildGradModel(Model forward, GradientRegistry registry) {
  return [forward_model = std::move(forward),
          grad_registry = std::move(registry)](
             AbstractContext* ctx,
             absl::Span<AbstractTensorHandle* const> inputs,
             absl::Span<AbstractTensorHandle*> outputs) -> Status {

      xla::PjRtClient::HostBufferSemantics::kImmutableOnlyDuringCall, nullptr,
      c_api_client->pjrt_c_client()->client->addressable_devices()[0]);
  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>();

    for (int j = 0; j < underlying_devices_.size(); ++j) {
      components.push_back(std::move(per_device_output_tensors[j][i]));
    }
    if (expected_output_shapes[i].IsFullyDefined()) {
      per_device_outputs.push_back(ParallelTensor::FromTensorHandles(
          *this, std::move(components),
          absl::Span<const int64_t>(expected_output_shapes[i].dim_sizes()),
          status));
    } else {