Home
last modified time | relevance | path

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

/aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tensorflow/transforms/
H A Dtpu_update_embedding_enqueue_op_inputs.cc66 func::FuncOp func_op, llvm::StringMap<Operation*>* enqueue_op_map, in FindTPUEmbeddingOps() argument
81 if (failed(ExtractEmbeddingAttribute(op, enqueue_op_map))) in FindTPUEmbeddingOps()
95 const llvm::StringMap<Operation*>& enqueue_op_map, in UpdateEmbeddingEnqueueOpInput() argument
99 for (const auto& it : enqueue_op_map) { in UpdateEmbeddingEnqueueOpInput()
151 llvm::StringMap<Operation*> enqueue_op_map; in runOnOperation() local
154 if (failed(FindTPUEmbeddingOps(func_op, &enqueue_op_map, in runOnOperation()
159 if (enqueue_op_map.size() != recv_activation_op_map.size()) { in runOnOperation()
167 if (failed(UpdateEmbeddingEnqueueOpInput(enqueue_op_map, in runOnOperation()