Searched refs:pjrt_client_ (Results 1 – 2 of 2) sorted by relevance
94 PjRtClient* pjrt_client() const { return pjrt_client_.get(); } in pjrt_client()95 std::shared_ptr<PjRtClient> shared_pjrt_client() { return pjrt_client_; } in shared_pjrt_client()98 return pjrt_client_->platform_name(); in platform_name()101 return pjrt_client_->addressable_device_count(); in addressable_device_count()103 int device_count() const { return pjrt_client_->device_count(); } in device_count()104 int task_id() const { return pjrt_client_->task_id(); } in task_id()119 return pjrt_client_->CreateChannelHandle(); in CreateChannelHandle()122 return pjrt_client_->CreateDeviceToHostChannelHandle(); in CreateDeviceToHostChannelHandle()125 return pjrt_client_->CreateHostToDeviceChannelHandle(); in CreateHostToDeviceChannelHandle()144 std::shared_ptr<PjRtClient> pjrt_client_; variable
34 : pjrt_client_(std::move(pjrt_client)) {} in PyClient()36 : pjrt_client_(std::move(pjrt_client)) {} in PyClient()40 auto span = pjrt_client_->devices(); in Devices()50 devices.reserve(pjrt_client_->addressable_devices().size()); in LocalDevices()51 for (PjRtDevice* device : pjrt_client_->addressable_devices()) { in LocalDevices()72 pjrt_client_->GetDefaultDeviceAssignment(num_replicas, num_partitions)); in GetDefaultDeviceAssignment()80 pjrt_client_->LookupDevice(device_id)); in GetDefaultDeviceAssignment()90 pjrt_client_->GetDefaultDeviceAssignment( in GetDefaultDeviceAssignment1D()96 pjrt_client_->LookupDevice(device_id)); in GetDefaultDeviceAssignment1D()106 TF_RET_CHECK(!pjrt_client_->addressable_devices().empty()); in PjRtBufferFromPyval()[all …]