// WithControlDependencies() would share the same NameMap with the parent.
  typedef std::unordered_map<string, int> NameMap;

  Impl(const std::shared_ptr<Graph>& graph,
       const std::shared_ptr<Status>& status,
       const std::shared_ptr<NameMap>& name_map,
       const std::shared_ptr<ShapeRefiner>& refiner);

  const string& name() const { return name_; }
  const std::vector<Operation>& control_deps() const { return control_deps_; }

class XlaDeviceContext : public DeviceContext {
 public:
  explicit XlaDeviceContext(
      std::shared_ptr<se::Stream> compute_stream,
      std::shared_ptr<se::Stream> host_to_device_stream,
      std::shared_ptr<se::Stream> device_to_host_stream,
      std::vector<std::shared_ptr<se::Stream>> device_to_device_streams,
      xla::LocalClient* client,

#include "tensorflow/core/lib/core/status.h"
#include "tsl/platform/errors.h"

namespace tensorflow {

bool XlaKernelCreator::CanCreateKernel(
    const FunctionLibraryRuntime& flr,
    const std::shared_ptr<const NodeProperties>& props) const {
  return CanCreateXlaKernel(props->node_def) &&
         !XlaOpRegistry::IsCompilationDevice(flr.device()->device_type());
}

static Status CreateXlaKernel(FunctionLibraryRuntime* flr,

  bool CanCreateKernel(
      const FunctionLibraryRuntime& flr,
      const std::shared_ptr<const NodeProperties>& props) const override;

  // Given a supported NodeDef, returns a XlaLaunchOp that computes the node.
  Status CreateKernel(FunctionLibraryRuntime* flr,
                      const std::shared_ptr<const NodeProperties>& props,
                      std::unique_ptr<OpKernel>* kernel) const override;
};

  explicit XlaPlatformInfo(
      const DeviceType device_type, se::Platform::Id platform_id,
      const XlaDevice::Metadata* xla_device_metadata,
      const PjRtBaseDevice::Metadata* pjrt_device_metadata,
      std::shared_ptr<se::DeviceMemoryAllocator> device_allocator)
      : device_type_(device_type),
        platform_id_(platform_id),
        xla_device_metadata_(xla_device_metadata),
        pjrt_device_metadata_(pjrt_device_metadata),

  // the RevivedObjects struct (which owns all functions). One alternative would
  // be to have RevivedObjects store shared_ptr<TFConcreteFunction> instead, and
  // change RestoredResource's constructor take shared_ptr<TFConcreteFunction>.
  // To keep things simple, I've stuck to raw pointers for now.
  //
  // Params:
  //  device - The device string associated with the SavedResource

  // streams.
  // It is legal to reset the definition event of a tensor when overwriting the
  // tensor's value (at which point, it is effectively a new tensor once again.)
  void ResetDefinitionEvent(std::shared_ptr<se::Event> event,
                            se::Stream* stream);

  // Refresh the status of streams_defined_on_. Return the first not-OK stream's
  // status or OK.
  Status RefreshStatusOfStreams();

class ClientSession::Impl {
 private:
  friend class ClientSession;

  Impl(Session* session, std::shared_ptr<Graph> graph)
      : session_(session), graph_(std::move(graph)) {}

  static SessionOptions MakeDefaultSessionOptions(const string& target);
  Status MaybeExtendGraph() const;

  std::unique_ptr<Session> session_;
  std::shared_ptr<Graph> graph_;

  mutable mutex mu_;
  mutable int last_num_graph_nodes_ TF_GUARDED_BY(mu_) = 0;

    return;
  }

  stream->WaitFor(definition_event_.get()).IgnoreError();
  streams_defined_on_.push_back(stream);
}

void XlaTensor::ResetDefinitionEvent(std::shared_ptr<se::Event> event,
                                     se::Stream* stream) {
  mutex_lock lock(mu_);
  definition_event_ = std::move(event);
  streams_defined_on_ = {stream};
}

  TF_FileStatistics base;
  int64_t generation_number;
} GcsFileStat;

typedef struct GCSFile {
  google::cloud::storage::Client gcs_client;  // owned
  bool compose;
  absl::Mutex block_cache_lock;
  std::shared_ptr<RamFileBlockCache> file_block_cache
      ABSL_GUARDED_BY(block_cache_lock);
  uint64_t block_size;  // Reads smaller than block_size will trigger a read
                        // of block_size.