[Auto Parallel]Add Embedding flops (#47978)

* c_embedding * add annotations * add annotations * revision * revise attrs
PaddlePaddle · Dec 1, 2022 · 08c5f4c · 08c5f4c
1 parent 9218e74
commit 08c5f4c
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git a/python/paddle/fluid/tests/unittests/test_profiler.py b/python/paddle/fluid/tests/unittests/test_profiler.py
@@ -297,6 +297,10 @@ def test_flops(self):
         self.assertTrue(
             flops('softmax', {'X': [[12, 12, 12]]}, {}) == 3 * 12 * 12 * 12
         )
+        self.assertTrue(
+            flops('c_embedding', {'Ids': [[12, 12]], 'W': [[12, 12, 3]]}, {})
+            == 0
+        )
 
 
 if __name__ == '__main__':

diff --git a/python/paddle/utils/flops.py b/python/paddle/utils/flops.py
@@ -60,6 +60,15 @@ def register(func):
     return register
 
 
+@register_flops("c_embedding")
+def _c_embedding_flops(input_shapes, attrs):
+    """FLOPs computation for c_embedding op.
+    For c_embedding(input):
+        equation: flops = 0
+    """
+    return 0
+
+
 @register_flops("dropout")
 def _dropout_flops(input_shapes, attrs):
     """FLOPs computation for dropout op.