Searched refs:dense_enqueue_op (Results 1 – 1 of 1) sorted by relevance
65 mlir::TF::DynamicEnqueueTPUEmbeddingArbitraryTensorBatchOp dense_enqueue_op = in ExpandOp() local69 mlir::OpBuilder builder(dense_enqueue_op); in ExpandOp()70 mlir::Location location = dense_enqueue_op->getLoc(); in ExpandOp()72 mlir::OperandRange feature = dense_enqueue_op.embedding_indices(); in ExpandOp()102 /*aggregation_weights=*/dense_enqueue_op.aggregation_weights(), in ExpandOp()104 dense_enqueue_op.mode_override(), in ExpandOp()105 /*device_ordinal=*/dense_enqueue_op.device_ordinal(), in ExpandOp()106 /*combiners=*/dense_enqueue_op.combiners()); in ExpandOp()107 dense_enqueue_op.erase(); in ExpandOp()