return false;
  }
  return stream_executor_->ClearAllocatorStats();
}

XlaDeviceContext::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,
    XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns,

      colocation_constraints_(),
      disable_shape_inference_(disable_shape_inference) {}

Scope::Impl::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)
    : graph_(graph),
      status_(status),
      name_map_(name_map),
      refiner_(refiner),

    TF_RETURN_IF_ERROR(EnsureStreamOkLocked(backend, "stream", &stream_,
                                            &need_new_device_context));
  }

  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;
  if (use_multiple_streams_) {
    TF_RETURN_IF_ERROR(EnsureStreamOkLocked(backend, "host_to_device_stream",

  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),

class TaggedValueHash;
using List = std::vector<TaggedValue>;
using ListPtr = std::shared_ptr<List>;
using Dict =
    absl::flat_hash_map<TaggedValue, TaggedValue, TaggedValueHash<TaggedValue>>;
using DictPtr = std::shared_ptr<Dict>;
using TuplePtr = std::shared_ptr<Tuple>;
using Func =
    std::function<absl::StatusOr<TaggedValue>(TaggedValue, TaggedValue)>;

  // If use_multiple_streams_, host to device transfers are performed using this
  // stream.
  std::shared_ptr<se::Stream> host_to_device_stream_ TF_GUARDED_BY(mu_);
  // If use_multiple_streams_, transfers between different devices are performed
  // using these streams.
  std::vector<std::shared_ptr<se::Stream>> device_to_device_streams_
      TF_GUARDED_BY(mu_);

  // See comments in options.

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;

      /*owned_args=*/{}, &non_donatable_input_indices));

  EXPECT_EQ(exec_args.size(), 2);

  std::shared_ptr<xla::Literal> literal1 = *exec_args[0]->ToLiteralSync();
  EXPECT_TRUE(xla::LiteralTestUtil::Equal(
      *literal1, xla::LiteralUtil::CreateR2<int32_t>({{1, 2, 3}})));

  std::shared_ptr<xla::Literal> literal2 = *exec_args[1]->ToLiteralSync();
  EXPECT_TRUE(xla::LiteralTestUtil::Equal(

    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};
}

  se::Platform::Id platform_id = nullptr;
  const XlaDevice::Metadata* xla_device_metadata = nullptr;
  const PjRtBaseDevice::Metadata* pjrt_device_metadata = nullptr;
  std::shared_ptr<se::DeviceMemoryAllocator> custom_allocator;

  const std::string& device_type = device_base->device_type();
  if (device_type == DEVICE_CPU) {
    platform_id = se::host::kHostPlatformId;