Searched refs:addressable_devices (Results 1 – 9 of 9) sorted by relevance
49 if (num_partitions == 1 && num_replicas <= addressable_devices().size()) { in GetDefaultDeviceAssignment()52 assignment(i, 0) = addressable_devices().at(i)->id(); in GetDefaultDeviceAssignment()77 std::vector<std::unique_ptr<LocalDeviceState>> addressable_devices; in BuildLocalDeviceStates() local81 addressable_devices.push_back(absl::make_unique<LocalDeviceState>( in BuildLocalDeviceStates()86 return std::move(addressable_devices); in BuildLocalDeviceStates()91 absl::Span<std::unique_ptr<LocalDeviceState> const> addressable_devices, in CreateBFCAllocator() argument93 CHECK_GT(addressable_devices.size(), 0); in CreateBFCAllocator()95 addressable_devices.front()->executor()->platform(); in CreateBFCAllocator()105 for (auto& local_device : addressable_devices) { in CreateBFCAllocator()148 absl::Span<std::unique_ptr<LocalDeviceState> const> addressable_devices) { in GetGpuDeviceAllocator() argument[all …]
160 virtual absl::Span<PjRtDevice* const> addressable_devices() const = 0;451 virtual absl::Span<PjRtDevice* const> addressable_devices() const = 0;
35 PjRtDevice* device = client->addressable_devices().at(0); in TEST()
140 absl::Span<PjRtDevice* const> addressable_devices() const override { in addressable_devices() function608 std::vector<PjRtDevice*> addressable_devices,642 absl::Span<PjRtDevice* const> addressable_devices() const override { in addressable_devices() function
1571 std::vector<PjRtDevice*> addressable_devices, in PjRtStreamExecutorExecutable() argument1578 addressable_devices_(std::move(addressable_devices)) { in PjRtStreamExecutorExecutable()2150 std::vector<PjRtDevice*> addressable_devices; in Compile() local2153 addressable_devices.reserve(num_replicas * num_partitions); in Compile()2166 addressable_devices.push_back(device); in Compile()2169 if (addressable_devices.empty()) { in Compile()2177 addressable_devices.front()->local_hardware_id()); in Compile()2188 std::move(addressable_devices), this); in Compile()
63 devices.reserve(executable_->addressable_devices().size()); in AddressableDevices()64 for (PjRtDevice* device : executable_->addressable_devices()) { in AddressableDevices()143 int num_computations = executable_->addressable_devices().size(); in ExecuteShardedOnLocalDevices()
50 devices.reserve(pjrt_client_->addressable_devices().size()); in LocalDevices()51 for (PjRtDevice* device : pjrt_client_->addressable_devices()) { in LocalDevices()106 TF_RET_CHECK(!pjrt_client_->addressable_devices().empty()); in PjRtBufferFromPyval()107 device = pjrt_client_->addressable_devices().front(); in PjRtBufferFromPyval()
234 for (auto device : client->addressable_devices()) { in OutfeedReceiverImpl()
1077 cache_entry->executable->pjrt_executable().addressable_devices().size(); in AddCacheEntry()