[Inductor] Fallback embedding when sparse is True (#150659)

leslie-fang-intel · pytorchmergebot · commit d6887f444fa6 · 2025-04-04T23:59:38.000Z
**Summary** Fix issue: #150656, fallback `embedding` when sparse is True. **Test Plan** ``` python -u -m pytest -s -v test/inductor/test_torchinductor.py -k test_embedding_sparse ``` Pull Request resolved: #150659 Approved by: https://github.com/jansel
diff --git a/test/inductor/test_torchinductor.py b/test/inductor/test_torchinductor.py
@@ -5360,6 +5360,19 @@ def test_embedding(self):
             (torch.randint(10, [2, 8]),),
         )
 
+    def test_embedding_sparse(self):
+        # Fix https://github.com/pytorch/pytorch/issues/150656
+        def fn(weight, indices):
+            return F.embedding(indices, weight, sparse=True)
+
+        indices = torch.randint(10, (2, 3))
+        weight = torch.randn(10, 3, requires_grad=True)
+
+        self.common(
+            fn,
+            (weight, indices),
+        )
+
     def test_mean(self):
         def fn(x):
             return (
diff --git a/test/inductor/test_torchinductor_codegen_dynamic_shapes.py b/test/inductor/test_torchinductor_codegen_dynamic_shapes.py
@@ -137,6 +137,7 @@ def run(*ex, **kwargs):
     "test_mul_index_expr_dynamic_shapes": TestFailure(("cpu",)),
     "test_flip_cat_dynamic_shapes": TestFailure(("cpu",)),
     "test_pad_single_dynamic_shapes": TestFailure(("cpu",)),
+    "test_embedding_sparse_dynamic_shapes": TestFailure(("cpu", "cuda", "xpu")),
     #
     # Failed to find for loop/triton kernel:
     #
diff --git a/torch/_inductor/lowering.py b/torch/_inductor/lowering.py
@@ -3376,6 +3376,11 @@ def fn(idx):
 
 @register_lowering(aten.embedding, type_promotion_kind=None)
 def embedding(weight, indices, padding_idx=-1, scale_grad_by_freq=False, sparse=False):
+    if sparse:
+        return fallback_handler(aten.embedding.default)(
+            weight, indices, padding_idx, scale_grad_by_freq, sparse
+        )
+
     assert not sparse
     assert isinstance(weight, TensorBox)
     assert isinstance(indices, TensorBox)

Original file line number	Diff line number	Diff line change
`@@ -137,6 +137,7 @@ def run(ex, *kwargs):`
`137`	`137`	`"test_mul_index_expr_dynamic_shapes": TestFailure(("cpu",)),`
`138`	`138`	`"test_flip_cat_dynamic_shapes": TestFailure(("cpu",)),`
`139`	`139`	`"test_pad_single_dynamic_shapes": TestFailure(("cpu",)),`
	`140`	`+ "test_embedding_sparse_dynamic_shapes": TestFailure(("cpu", "cuda", "xpu")),`
`140`	`141`	`#`
`141`	`142`	`# Failed to find for loop/triton kernel:`
`142`	`143`	`#`