Home
last modified time | relevance | path

Searched refs:executables_ (Results 1 – 9 of 9) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/python/tpu_driver/client/
Dtpu_client.h302 for (auto it = executables_.begin(); it != executables_.end(); ++it) { in ~PyTpuExecutable()
318 CHECK_GE(executables_.size(), 1); in SizeOfGeneratedCodeInBytes()
319 return executables_.begin()->second->size_in_bytes(); in SizeOfGeneratedCodeInBytes()
352 void Delete() { executables_.clear(); } in Delete()
366 std::map<int, std::unique_ptr<tpu_driver::LoadedProgramHandle>> executables_; variable
Dtpu_client.cc530 CHECK(executables_.find(replica) == executables_.end()) in PyTpuExecutable()
532 executables_[replica] = in PyTpuExecutable()
539 CHECK_LE(executables_.size(), client_->device_count()); in PyTpuExecutable()
585 executables_.find(replica)->second.get(), inputs, in ExecuteHelper()
/external/tensorflow/tensorflow/compiler/xla/python/
Dpy_executable.cc34 next_ = client_->executables_; in PyExecutable()
35 client_->executables_ = this; in PyExecutable()
50 if (client_->executables_ == this) { in ~PyExecutable()
51 client_->executables_ = next_; in ~PyExecutable()
Dpmap_lib.cc144 for (const auto& entry : executables_) { in ~PmapFunction()
162 int cache_size() const { return executables_.size(); } in cache_size()
183 executables_; member in jax::PmapFunction
194 auto found_iterator = executables_.find(signature); in GetCacheEntryIfPresent()
195 if (found_iterator != executables_.end()) { // Cache hit! in GetCacheEntryIfPresent()
215 executables_.emplace(signature, std::make_unique<PmapCacheEntry>()); in AddCacheEntry()
Dpy_client.h150 PyExecutable* executables_ = nullptr; variable
Djax_jit.cc863 int cache_size() const { return executables_.size(); } in cache_size()
884 absl::flat_hash_map<CallSignature, std::unique_ptr<CacheEntry>> executables_; member in jax::__anone04173081411::CompiledFunction
931 for (const auto& entry : executables_) { in ~CompiledFunction()
1028 auto found_iterator = executables_.find(signature); in GetCacheEntryIfPresent()
1029 if (found_iterator != executables_.end()) { // Cache hit! in GetCacheEntryIfPresent()
1048 auto result = executables_.emplace(signature, std::make_unique<CacheEntry>()); in AddCacheEntry()
Dpy_client.cc280 for (PyExecutable* executable = executables_; executable; in HeapProfile()
/external/tensorflow/tensorflow/compiler/xla/pjrt/
Dpjrt_stream_executor_client.h618 return executables_[0]->build_options().num_replicas(); in num_replicas()
622 return executables_[0]->build_options().num_partitions(); in num_partitions()
627 for (auto& executable : executables_) { in SizeOfGeneratedCodeInBytes()
662 void Delete() override { executables_.clear(); } in Delete()
665 return executables_; in executables()
711 std::vector<std::shared_ptr<LocalExecutable>> executables_; variable
Dpjrt_stream_executor_client.cc1579 executables_.reserve(executables.size()); in PjRtStreamExecutorExecutable()
1581 executables_.emplace_back(std::move(executable)); in PjRtStreamExecutorExecutable()
1601 if (executables_.size() > 1) { in PjRtStreamExecutorExecutable()
1602 CHECK_EQ(num_partitions, executables_.size()) in PjRtStreamExecutorExecutable()
1603 << "Number of executables " << executables_.size() in PjRtStreamExecutorExecutable()
1609 parameters_that_must_be_donated_.reserve(executables_.size()); in SetUpDonation()
1610 for (auto& executable : executables_) { in SetUpDonation()
1621 Executable* executable = executables_[0]->executable(); in name()
1772 executables_[executable_idx]->RunAsync(std::move(execution_inputs), in EnqueueExecution()
1801 [references{std::make_tuple(executables_[executable_idx], in EnqueueExecution()
[all …]