Add impl for aten.embedding and lowering.

PiperOrigin-RevId: 773846228

Author: junjiang-lab

ai_edge_torch/odml_torch/experimental/torch_tfl/_decomps.py

@@ -283,6 +283,12 @@ def _aten_where_self_decomp(condition, x, y):
   return torch.ops.tfl.select_v2(condition, x, y)
 
 
+@register_decomp(torch.ops.aten.embedding.default)
+def _aten_embedding_decomp(weight, indices, padding_idx=-1):
+  # TODO: b/425747317 - Decomp to tfl.embedding_lookup once it's ready.
+  return torch.ops.tfl.gather(weight, indices, axis=0)
+
+
 @register_decomp(torch.ops.aten._softmax.default)
 def _aten__softmax_decomp(
     x, dim: int, half_to_float: bool  # pylint: disable=unused-argument

ai_edge_torch/odml_torch/experimental/torch_tfl/_lowerings.py

@@ -600,6 +600,40 @@ def _tfl_select_v2_lowering(
   )
 
 
+@lower(torch.ops.tfl.embedding_lookup.default)
+def _tfl_embedding_lookup_lowering(
+    lctx: LoweringContext,
+    indices: ir.Value,
+    weight: ir.Value,
+) -> ir.Value:
+  return _ir_operation(
+      "tfl.embedding_lookup",
+      results=lowering_utils.node_meta_to_ir_types(lctx.node),
+      operands=[indices, weight],
+  )
+
+
+@lower(torch.ops.tfl.gather.default)
+def _tfl_gather_lowering(
+    lctx: LoweringContext,
+    x: ir.Value,
+    indices: ir.Value,
+    axis: int,
+    batch_dims: int = 0,
+) -> ir.Value:
+  return _ir_operation(
+      "tfl.gather",
+      results=lowering_utils.node_meta_to_ir_types(lctx.node),
+      operands=[x, indices],
+      attributes={
+          "axis": ir.IntegerAttr.get(ir.IntegerType.get_signless(32), axis),
+          "batch_dims": ir.IntegerAttr.get(
+              ir.IntegerType.get_signless(32), batch_dims
+          ),
+      },
+  )
+
+
 @lower(torch.ops.tfl.softmax.default)
 def _tfl_softmax_lowering(
     lctx: LoweringContext,

ai_edge_torch/odml_torch/experimental/torch_tfl/_ops.py

@@ -273,6 +273,20 @@ def tfl_select_v2(
   return torch.where(condition, x, y)
 
 
+@custom_op_with_fake("tfl::embedding_lookup")
+def tfl_embedding_lookup(
+    indices: torch.Tensor, weight: torch.Tensor
+) -> torch.Tensor:
+  return torch.nn.functional.embedding(indices, weight)
+
+
+@custom_op_with_fake("tfl::gather")
+def tfl_gather(
+    input: torch.Tensor, indices: torch.Tensor, axis: int
+) -> torch.Tensor:
+  return torch.index_select(input, axis, indices)
+
+
 @custom_op_with_fake("tfl::softmax")
 def tfl_softmax(x: torch.Tensor) -> torch.Tensor:
   return torch.nn.functional.softmax(x, dim=-1)

ai_edge_torch/odml_torch/experimental/torch_tfl/test/test_torch_tfl_impls.py

@@ -186,6 +186,7 @@ def _assert_export_and_close(
       ("aten_select_int_0", torch.ops.aten.select.int, (rnd(torch.float32, (2, 3, 4)), 0, 1,), dict()),
       ("aten_select_int_1", torch.ops.aten.select.int, (rnd(torch.float32, (2, 3, 4)), 1, 1,), dict()),
       ("aten_where_self_0", torch.ops.aten.where.self, (rnd(torch.bool, (10, 10)), rnd(torch.float32, (10, 10)), rnd(torch.float32, (10, 10)),), dict()),
+      ("aten_embedding_0", torch.ops.aten.embedding.default, (rnd(torch.float32, (10, 10)), torch.tensor([0, 2, 4, 6, 8]),), dict()),
       ("aten__softmax_0", torch.ops.aten._softmax.default, (rnd(torch.float32, (10, 10)), -1, False), dict()),
       ("aten__softmax_1", torch.ops.aten._softmax.default, (rnd(torch.float32, (1, 10)), -1, False), dict()),
       ("aten__softmax_2", torch.ops.aten._softmax.default, (rnd(torch.float32, (10, 10)), 0, False), dict()),