Home
last modified time | relevance | path

Searched refs:GetDefaultDeviceAssignment (Results 1 – 15 of 15) sorted by relevance

/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/pjrt/
H A Dgpu_device.cc120 xla::StatusOr<xla::DeviceAssignment> GetDefaultDeviceAssignment(
138 xla::StatusOr<xla::DeviceAssignment> GpuClient::GetDefaultDeviceAssignment( in GetDefaultDeviceAssignment() function in xla::__anonb78c89a70111::GpuClient
148 return PjRtStreamExecutorClient::GetDefaultDeviceAssignment(num_replicas, in GetDefaultDeviceAssignment()
H A Dtpu_client.cc120 StatusOr<DeviceAssignment> PjRtTpuClient::GetDefaultDeviceAssignment( in GetDefaultDeviceAssignment() function in xla::PjRtTpuClient
131 return PjRtStreamExecutorClient::GetDefaultDeviceAssignment(num_replicas, in GetDefaultDeviceAssignment()
H A Dpjrt_c_api_client.h145 StatusOr<DeviceAssignment> GetDefaultDeviceAssignment( in GetDefaultDeviceAssignment() function
148 return wrapped_->GetDefaultDeviceAssignment(num_replicas, num_partitions); in GetDefaultDeviceAssignment()
H A Dtpu_client.h81 StatusOr<DeviceAssignment> GetDefaultDeviceAssignment(
H A Dpjrt_client.h480 virtual StatusOr<DeviceAssignment> GetDefaultDeviceAssignment(
489 virtual StatusOr<DeviceAssignment> GetDefaultDeviceAssignment( in GetDefaultDeviceAssignment() function
H A Dtfrt_cpu_pjrt_client.h149 StatusOr<DeviceAssignment> GetDefaultDeviceAssignment(
H A Dpjrt_stream_executor_client.h190 StatusOr<DeviceAssignment> GetDefaultDeviceAssignment(
H A Dtfrt_cpu_pjrt_client.cc232 StatusOr<DeviceAssignment> TfrtCpuClient::GetDefaultDeviceAssignment( in GetDefaultDeviceAssignment() function in xla::TfrtCpuClient
343 return this->GetDefaultDeviceAssignment(num_replicas, num_partitions); in Compile()
H A Dpjrt_stream_executor_client.cc257 StatusOr<DeviceAssignment> PjRtStreamExecutorClient::GetDefaultDeviceAssignment( in GetDefaultDeviceAssignment() function in xla::PjRtStreamExecutorClient
2410 return this->GetDefaultDeviceAssignment(num_replicas, num_partitions); in GetExecutableExtras()
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/python/tpu_driver/client/
H A Dtpu_client_extension.cc52 client->GetDefaultDeviceAssignment( in PYBIND11_MODULE()
72 client->GetDefaultDeviceAssignment( in PYBIND11_MODULE()
H A Dtpu_client.cc144 StatusOr<DeviceAssignment> PyTpuClient::GetDefaultDeviceAssignment( in GetDefaultDeviceAssignment() function in xla::PyTpuClient
839 client->GetDefaultDeviceAssignment( in Compile()
H A Dtpu_client.h131 virtual StatusOr<DeviceAssignment> GetDefaultDeviceAssignment(
/aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/python/
H A Dpy_client.cc181 PyClient::GetDefaultDeviceAssignment(int num_replicas, int num_partitions) { in GetDefaultDeviceAssignment() function in xla::PyClient
184 pjrt_client_->GetDefaultDeviceAssignment(num_replicas, num_partitions)); in GetDefaultDeviceAssignment()
202 pjrt_client_->GetDefaultDeviceAssignment( in GetDefaultDeviceAssignment1D()
H A Dpy_client.h133 GetDefaultDeviceAssignment(int num_replicas, int num_partitions);
H A Dxla.cc214 &PyClient::GetDefaultDeviceAssignment) in PYBIND11_MODULE()