masahi edited a comment on pull request #7149: URL: https://github.com/apache/tvm/pull/7149#issuecomment-773134472
I'm interested in this discussion. I'm looking at how to support PyTorch [EmbeddingBag](https://pytorch.org/docs/stable/generated/torch.nn.EmbeddingBag.html) op, which is used in Facebook DLRM model and which seems very similar to TF sparse segment sum op this PR is adding. The same ops exist in caffe2, too https://caffe2.ai/docs/sparse-operations.html So I'd say this op is not specific to TF and generally useful for recsys or possibly NLP use cases. This op can be implemented by composing embedding op + reduce, and that is how ONNX exports PyTorch `EmbeddingBag` op. But that would involve materializing the embedding which is huge for DLRM. Both PyTorch and caffe2 use custom fused embedding lookup + on the fly reduce to efficiently implement these ops. We cannot rely on automatic fusion of embedding + reduce, because our op fusion doesn't fuse any ops before a reduction op (if I remember correctly) cc @jwfromm @mbrookhart ---------------------------------------------------------------- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. For queries about this service, please contact Infrastructure at: [email protected]
