Searched refs:TpuProgramGroup (Results 1 – 8 of 8) sorted by relevance
36 TPUExecutableInfoProto TpuProgramGroup::ConstructExecutableInfo( in ConstructExecutableInfo()52 TPUHostTransferInfoProto TpuProgramGroup::ConstructHostTransferInfo( in ConstructHostTransferInfo()68 xla::HloProto TpuProgramGroup::ConstructHloMetadata( in ConstructHloMetadata()84 void TpuProgramGroup::Initialize( in Initialize()114 bool TpuProgramGroup::has_sharding_program() const { in has_sharding_program()123 size_t TpuProgramGroup::program_count() const { return tpu_programs_.size(); } in program_count()125 int64_t TpuProgramGroup::program_size() const { in program_size()133 bool TpuProgramGroup::LogProgramMemorySummary() { in LogProgramMemorySummary()141 void TpuProgramGroup::UnloadAndDestroyPrograms() { in UnloadAndDestroyPrograms()153 TpuProgramGroup::TpuProgramGroup(TpuProgramGroup&& other) in TpuProgramGroup() function in tensorflow::tpu::TpuProgramGroup[all …]
87 class TpuProgramGroup : public TpuProgramGroupInterface {107 TpuProgramGroup() = default;108 TpuProgramGroup(TpuProgramGroup&& other);109 TpuProgramGroup& operator=(TpuProgramGroup&&) = delete;184 TF_DISALLOW_COPY_AND_ASSIGN(TpuProgramGroup);
44 TpuProgramGroup tpu_program_group) { in PopulateEntry()51 absl::make_unique<TpuProgramGroup>(std::move(tpu_program_group)); in PopulateEntry()64 entry->tpu_program_group = absl::make_unique<TpuProgramGroup>(); in CreateAndInitializeCompiledSubgraph()96 TpuProgramGroup tpu_program_group; in InitializeEntry()118 TpuProgramGroup sharding_programs; in InitializeEntry()126 TpuProgramGroup unsharding_programs; in InitializeEntry()
54 auto tpu_program_group = absl::make_unique<TpuProgramGroup>(); in DeserializeRpcResponseToCacheEntry()82 const TpuProgramGroup* tpu_program_group = in SerializeCacheEntryToBufferSlices()83 tensorflow::down_cast<const TpuProgramGroup*>( in SerializeCacheEntryToBufferSlices()
37 return TpuProgramGroup::CompileAndBuild(compilation_request, mesh_state, in Compile()
648 const tpu::TpuProgramGroup* tpu_program_group = in DoWork()649 tensorflow::down_cast<const tpu::TpuProgramGroup*>( in DoWork()
125 const tpu::TpuProgramGroup* tpu_program_group, xla::Backend* const backend, in GetExecutableAliasConfig()191 const uint32 rng_seed, const tpu::TpuProgramGroup* tpu_program_group, in RunExecutable()233 const uint32 rng_seed, const tpu::TpuProgramGroup* tpu_program_group, in ExecuteTPUProgram()313 const tpu::TpuProgramGroup* tpu_program_group = in DoWork()314 tensorflow::down_cast<const tpu::TpuProgramGroup*>( in DoWork()436 const tpu::TpuProgramGroup* tpu_program_group = in DoWork()437 tensorflow::down_cast<const tpu::TpuProgramGroup*>( in DoWork()
86 return tensorflow::tpu::TpuProgramGroup::CompileAndBuild( in Compile()