ai-edge-torch-nightly 0.2.0.dev20240730__py3-none-any.whl → 0.2.0.dev20240731__py3-none-any.whl

ai_edge_torch/convert/fx_passes/build_aten_composite_pass.py

@@ -223,21 +223,25 @@ def _aten_embedding(gm: GraphModule, node: Node):
     full_kwargs = args_mapper.get_full_kwargs(args, kwargs)
     _, embedding_dim = full_kwargs["weight"].size()
     idx = full_kwargs["indices"]
-    idx = idx.type(torch.int)
-    B, T = idx.size()
-
-    idx = torch.reshape(idx, (B * T,))
-
-    builder = StableHLOCompositeBuilder("odml.embedding_lookup")
-    full_kwargs["indices"], full_kwargs["weight"] = builder.mark_inputs(
-        idx,
-        full_kwargs["weight"],
-    )
-    output = op(**full_kwargs)
-    output = builder.mark_outputs(output)
-
-    output = torch.reshape(output, (B, T, embedding_dim))
-    return output
+    # TODO(b/356458830): Handle relative positional encoding
+    if len(idx.size()) == 2:
+      idx = idx.type(torch.int)
+      B, T = idx.size()
+
+      idx = torch.reshape(idx, (B * T,))
+
+      builder = StableHLOCompositeBuilder("odml.embedding_lookup")
+      full_kwargs["indices"], full_kwargs["weight"] = builder.mark_inputs(
+          idx,
+          full_kwargs["weight"],
+      )
+      output = op(**full_kwargs)
+      output = builder.mark_outputs(output)
+
+      output = torch.reshape(output, (B, T, embedding_dim))
+      return output
+    else:
+      return op(**full_kwargs)
 
   node.target = embedding
 

{ai_edge_torch_nightly-0.2.0.dev20240730.dist-info → ai_edge_torch_nightly-0.2.0.dev20240731.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.2.0.dev20240730
+Version: 0.2.0.dev20240731
 Summary: Supporting PyTorch models with the Google AI Edge TFLite runtime.
 Home-page: https://github.com/google-ai-edge/ai-edge-torch
 Keywords: On-Device ML,AI,Google,TFLite,PyTorch,LLMs,GenAI

{ai_edge_torch_nightly-0.2.0.dev20240730.dist-info → ai_edge_torch_nightly-0.2.0.dev20240731.dist-info}/RECORD

@@ -7,7 +7,7 @@ ai_edge_torch/convert/converter.py,sha256=hSrW6A-kix9cjdD6CuLL7rseWrLKoV6GRy-iUS
 ai_edge_torch/convert/to_channel_last_io.py,sha256=zo5tY3yDhY_EPCkrL1XSXs2uRFS8B4_qu08dSjNsUGk,2778
 ai_edge_torch/convert/fx_passes/__init__.py,sha256=EPs4PSIDLuRH5EBETi6deaOvaaf_Q4xD3_9NVcR7x8o,2810
 ai_edge_torch/convert/fx_passes/_pass_base.py,sha256=ijVyDclPnd6a0DWWUJkwR4igj6f82S-cE1-83QGPvgw,1652
-ai_edge_torch/convert/fx_passes/build_aten_composite_pass.py,sha256=ouV1CD_t5-MpDgr-7_zUG6vPrRYDT3-YWq81oZqCi9M,7924
+ai_edge_torch/convert/fx_passes/build_aten_composite_pass.py,sha256=NwJp8GbZWRtQML7hG2q5sS1W92RKiRsgDDdMLy4uIBc,8079
 ai_edge_torch/convert/fx_passes/build_interpolate_composite_pass.py,sha256=6m_vcycd9f3OQgQLx2hhQjsKfOqdxE5EkjzqrxqyAQM,4168
 ai_edge_torch/convert/fx_passes/canonicalize_pass.py,sha256=UX6dJsxCqSkftXXvNBV-i7Bjk6H7qTyqzUnE640Itfg,1673
 ai_edge_torch/convert/fx_passes/inject_mlir_debuginfo_pass.py,sha256=aRT8hTS3n9ie28lgu6mygtFO6Ypwu0qjNb0c81v9HLs,2448
@@ -125,8 +125,8 @@ ai_edge_torch/quantize/quant_config.py,sha256=eO9Ra160ITjQSyRBEGy6nNIVH3gYacSWDd
 ai_edge_torch/testing/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/testing/model_coverage/__init__.py,sha256=5P8J6Zk5YYtDvTBucFvB9NGSRI7Gw_24WnrbhXgycEE,765
 ai_edge_torch/testing/model_coverage/model_coverage.py,sha256=kzIulTldq8R9E-lAZsvfSTvLu3FYEX7b9DyYM3qisXM,4485
-ai_edge_torch_nightly-0.2.0.dev20240730.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.2.0.dev20240730.dist-info/METADATA,sha256=f9fGXKPiQpY75w77AU-rJxkSdq6WJPJ9Jvu71s7IIhk,1889
-ai_edge_torch_nightly-0.2.0.dev20240730.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-ai_edge_torch_nightly-0.2.0.dev20240730.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.2.0.dev20240730.dist-info/RECORD,,
+ai_edge_torch_nightly-0.2.0.dev20240731.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.2.0.dev20240731.dist-info/METADATA,sha256=B2Nf7g2PWOU-bYTAByfDNV_FAKy3ah88O-Plsk-uW_M,1889
+ai_edge_torch_nightly-0.2.0.dev20240731.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+ai_edge_torch_nightly-0.2.0.dev20240731.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.2.0.dev20240731.dist-info/RECORD,,