Searched refs:PjRtStreamExecutorClient (Results 1 – 7 of 7) sorted by relevance
121 class PjRtStreamExecutorClient : public PjRtClient {124 explicit PjRtStreamExecutorClient(131 ~PjRtStreamExecutorClient() override = default;419 friend class PjRtStreamExecutorClient; variable473 PjRtStreamExecutorClient* client() const override { return client_; } in client()587 PjRtStreamExecutorClient* const client_;609 PjRtStreamExecutorClient* client);613 PjRtStreamExecutorClient* client() const override { return client_; } in client()674 friend class PjRtStreamExecutorClient;709 PjRtStreamExecutorClient* const client_;
203 PjRtStreamExecutorClient::PjRtStreamExecutorClient( in PjRtStreamExecutorClient() function in xla::PjRtStreamExecutorClient254 StatusOr<DeviceAssignment> PjRtStreamExecutorClient::GetDefaultDeviceAssignment( in GetDefaultDeviceAssignment()261 PjRtStreamExecutorClient::GetHloCostAnalysis() { in GetHloCostAnalysis()372 auto* se_client = tensorflow::down_cast<PjRtStreamExecutorClient*>(client); in AllocateDestinationBuffer()647 PjRtStreamExecutorClient::BufferFromHostBuffer( in BufferFromHostBuffer()828 PjRtStreamExecutorClient::CreateUninitializedBuffer(const Shape& shape, in CreateUninitializedBuffer()834 PjRtStreamExecutorClient::CreateUninitializedBuffer( in CreateUninitializedBuffer()859 PjRtStreamExecutorClient::BufferFromHostLiteral(const LiteralSlice& literal, in BufferFromHostLiteral()920 void PjRtStreamExecutorClient::MakeCrossHostReceiveBuffers( in MakeCrossHostReceiveBuffers()958 PjRtStreamExecutorClient::CreateViewOfDeviceBuffer( in CreateViewOfDeviceBuffer()[all …]
76 class PjRtTpuClient : public PjRtStreamExecutorClient {94 : PjRtStreamExecutorClient(kTpuName, client, std::move(devices), task_id, in PjRtTpuClient()111 return PjRtStreamExecutorClient::GetDefaultDeviceAssignment(num_replicas, in GetDefaultDeviceAssignment()
39 class GpuClient : public xla::PjRtStreamExecutorClient {41 using xla::PjRtStreamExecutorClient::PjRtStreamExecutorClient;57 return PjRtStreamExecutorClient::GetDefaultDeviceAssignment(num_replicas, in GetDefaultDeviceAssignment()
54 return std::unique_ptr<PjRtClient>(std::make_unique<PjRtStreamExecutorClient>( in GetInterpreterClient()
60 return std::unique_ptr<PjRtClient>(std::make_unique<PjRtStreamExecutorClient>( in GetCpuClient()
100 return std::unique_ptr<PjRtClient>(std::make_unique<PjRtStreamExecutorClient>( in GetCpuClientWithNonLocalDevice()