[fx] refactored the file structure of patched function and module (#1238)

* [fx] refactored the file structure of patched function and module * polish code
2025-09-23 10:30:03 +00:00 · 2022-07-12 15:01:58 +08:00
parent 17ed33350b
commit 7531c6271f
15 changed files with 353 additions and 318 deletions
--- a/colossalai/fx/tracer/meta_patch/patched_function/embedding.py
+++ b/colossalai/fx/tracer/meta_patch/patched_function/embedding.py
@@ -0,0 +1,13 @@
+import torch
+from ..registry import meta_patched_function
+
+
+@meta_patched_function.register(torch.nn.functional.embedding)
+def torch_nn_functional_embedding(input,
+                                  weight,
+                                  padding_idx=None,
+                                  max_norm=None,
+                                  norm_type=2.0,
+                                  scale_grad_by_freq=False,
+                                  sparse=False):
+    return torch.empty(*input.shape, weight.shape[-1], device="meta")