Searched refs:EmbeddingLookUpPSKernel (Results 1 – 3 of 3) sorted by relevance
31 void EmbeddingLookUpPSKernel::InitKernel( in InitKernel()70 void EmbeddingLookUpPSKernel::ReInit(const std::vector<std::vector<size_t>> &shapes) { in ReInit()85 bool EmbeddingLookUpPSKernel::Execute(const std::vector<AddressPtr> &inputs, const std::vector<Addr… in Execute()90 void EmbeddingLookUpPSKernel::UpdateEmbeddings(float *embedding_table, const size_t *lookup_ids, in UpdateEmbeddings()107 const std::vector<size_t> &EmbeddingLookUpPSKernel::input_sizes() const { return input_shape_; } in input_sizes()109 const std::vector<size_t> &EmbeddingLookUpPSKernel::output_sizes() const { return GetOutputSizeList… in output_sizes()111 const std::vector<size_t> &EmbeddingLookUpPSKernel::workspace_sizes() const { return GetWorkspaceSi… in workspace_sizes()
28 class EmbeddingLookUpPSKernel : public EmbeddingLookUpCPUKernel, public PServerKernel {30 EmbeddingLookUpPSKernel(size_t rank_id, size_t pserver_num, size_t worker_num) in EmbeddingLookUpPSKernel() function32 ~EmbeddingLookUpPSKernel() override = default;
222 …std::make_shared<kernel::ps::EmbeddingLookUpPSKernel>(server_node_->rank_id(), pserver_num_, worke… in InitEmbeddingTable()