Home
last modified time | relevance | path

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

/external/tensorflow/tensorflow/compiler/xla/pjrt/
Dtpu_client.cc82 StatusOr<DeviceAssignment> GetDefaultDeviceAssignment(
100 StatusOr<DeviceAssignment> PjRtTpuClient::GetDefaultDeviceAssignment( in GetDefaultDeviceAssignment() function in xla::__anon240436b40111::PjRtTpuClient
111 return PjRtStreamExecutorClient::GetDefaultDeviceAssignment(num_replicas, in GetDefaultDeviceAssignment()
Dgpu_device.cc43 xla::StatusOr<xla::DeviceAssignment> GetDefaultDeviceAssignment(
47 xla::StatusOr<xla::DeviceAssignment> GpuClient::GetDefaultDeviceAssignment( in GetDefaultDeviceAssignment() function in xla::__anon0c1acb180111::GpuClient
57 return PjRtStreamExecutorClient::GetDefaultDeviceAssignment(num_replicas, in GetDefaultDeviceAssignment()
Dpjrt_client.h178 virtual StatusOr<DeviceAssignment> GetDefaultDeviceAssignment(
Dpjrt_stream_executor_client.h164 StatusOr<DeviceAssignment> GetDefaultDeviceAssignment(
Dpjrt_stream_executor_client.cc254 StatusOr<DeviceAssignment> PjRtStreamExecutorClient::GetDefaultDeviceAssignment( in GetDefaultDeviceAssignment() function in xla::PjRtStreamExecutorClient
2133 return this->GetDefaultDeviceAssignment(num_replicas, num_partitions); in Compile()
/external/tensorflow/tensorflow/compiler/xla/python/
Dpy_client.cc69 PyClient::GetDefaultDeviceAssignment(int num_replicas, int num_partitions) { in GetDefaultDeviceAssignment() function in xla::PyClient
72 pjrt_client_->GetDefaultDeviceAssignment(num_replicas, num_partitions)); in GetDefaultDeviceAssignment()
90 pjrt_client_->GetDefaultDeviceAssignment( in GetDefaultDeviceAssignment1D()
Dpy_client.h112 GetDefaultDeviceAssignment(int num_replicas, int num_partitions);
Dxla.cc213 &PyClient::GetDefaultDeviceAssignment) in PYBIND11_MODULE()
/external/tensorflow/tensorflow/compiler/xla/python/tpu_driver/client/
Dtpu_client_extension.cc42 client->GetDefaultDeviceAssignment( in PYBIND11_MODULE()
62 client->GetDefaultDeviceAssignment( in PYBIND11_MODULE()
Dtpu_client.h106 virtual StatusOr<DeviceAssignment> GetDefaultDeviceAssignment(
Dtpu_client.cc133 StatusOr<DeviceAssignment> PyTpuClient::GetDefaultDeviceAssignment( in GetDefaultDeviceAssignment() function in xla::PyTpuClient
829 client->GetDefaultDeviceAssignment( in Compile()