ai-edge-torch-nightly 0.2.0.dev20240802__py3-none-any.whl → 0.2.0.dev20240803__py3-none-any.whl

ai_edge_torch/__init__.py

@@ -17,6 +17,7 @@ from .convert.converter import convert
 from .convert.converter import signature
 from .convert.to_channel_last_io import to_channel_last_io
 from .model import Model
+from .version import __version__
 
 
 def load(path: str) -> Model:

ai_edge_torch/convert/fx_passes/build_aten_composite_pass.py

@@ -15,6 +15,7 @@
 
 import copy
 import functools
+from functools import reduce
 from typing import Any, Callable
 
 from ai_edge_torch.hlfb import StableHLOCompositeBuilder
@@ -228,25 +229,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"]
-    # 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)
+
+    # Explicitly cast to INT32. This places the CastOp outside of the HLFB.
+    idx = idx.type(torch.int)
+    original_idx_shape = idx.size()
+
+    # Explicitly reshape to 1D. This places the ReshapeOp outside of the HLFB.
+    idx = torch.reshape(idx, (idx.numel(),))
+
+    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)
+
+    # Explicitly reshape back to the original shape. This places the ReshapeOp outside of the HLFB.
+    output = torch.reshape(output, (*(original_idx_shape), embedding_dim))
+    return output
 
   node.target = embedding
 

ai_edge_torch/version.py

@@ -0,0 +1,16 @@
+# Copyright 2024 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+__version__ = "0.2.0.dev20240803"

{ai_edge_torch_nightly-0.2.0.dev20240802.dist-info → ai_edge_torch_nightly-0.2.0.dev20240803.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.2.0.dev20240802
+Version: 0.2.0.dev20240803
 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.dev20240802.dist-info → ai_edge_torch_nightly-0.2.0.dev20240803.dist-info}/RECORD

@@ -1,5 +1,6 @@
-ai_edge_torch/__init__.py,sha256=CNDboRP4zQBpz2hznNCQWcQCARvNXUm3DMa1Dw_XXFg,1067
+ai_edge_torch/__init__.py,sha256=WTuorXzCALfr89FC4kX_PBtKOQLipN1hcW2tMDSQW9w,1100
 ai_edge_torch/model.py,sha256=pSyY9O7J1i-SJu7g4mFD853MJBNFE6LSzBgJw7dtWuI,4494
+ai_edge_torch/version.py,sha256=JKJ2vopU7w5KwFlGOP8TesbgJffKwyyVUfrJTUxVDR8,706
 ai_edge_torch/convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/convert/conversion.py,sha256=bkOyaTTZR9lT1VJMxwCSjcplheYv1HNSwt8A9kEo388,4183
 ai_edge_torch/convert/conversion_utils.py,sha256=GAOFepARe_vxOaetplMBBaexxojSijJzXvkxft88-Lc,13945
@@ -7,7 +8,7 @@ ai_edge_torch/convert/converter.py,sha256=6BoHl_GEIOkTr1oBg-VzZb5tr6Rv9yDwxKczYd
 ai_edge_torch/convert/to_channel_last_io.py,sha256=b7Q0_6Lam6IV-3TyhabVTMS7j0ppFpKDOIHTNAw2PnI,2814
 ai_edge_torch/convert/fx_passes/__init__.py,sha256=D4Xe8YmeP2N0yEN_bc7pEJH47KkwGFf4COZOILmDL4w,2809
 ai_edge_torch/convert/fx_passes/_pass_base.py,sha256=WVYZuocpygHAzk9u1GNoGowAIOHTlJXyA_NklmYkRms,1672
-ai_edge_torch/convert/fx_passes/build_aten_composite_pass.py,sha256=DGoKclQuxjZChGJCxKs-07zufcFxrIzKBS7Ymi-lPiQ,8079
+ai_edge_torch/convert/fx_passes/build_aten_composite_pass.py,sha256=QaZ5JV7RazGbC2Khdai795vlO5jDc3yhgx3HHNmzHDs,8246
 ai_edge_torch/convert/fx_passes/build_interpolate_composite_pass.py,sha256=BWSU9nkD5DzxHI_WGcs9uH6qKWCw0XB2etDEV6PsZkg,4181
 ai_edge_torch/convert/fx_passes/canonicalize_pass.py,sha256=eW0Yae2cL2ALYVkhsuk3wX8v41P6bkGaABtRgdPCdxk,1672
 ai_edge_torch/convert/fx_passes/inject_mlir_debuginfo_pass.py,sha256=aRT8hTS3n9ie28lgu6mygtFO6Ypwu0qjNb0c81v9HLs,2448
@@ -125,8 +126,8 @@ ai_edge_torch/quantize/quant_config.py,sha256=yP93mRbsB03K1_dYCRIKgxRNEP4EJOYF68
 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=049yZFfnlVefQJAXkcn84ETzVneaZIlz8e0X1BW3vvI,4520
-ai_edge_torch_nightly-0.2.0.dev20240802.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.2.0.dev20240802.dist-info/METADATA,sha256=Bng_BviZH6NODVQolxehLCzUIDv6i6cVDB5Ddfj-uhc,1889
-ai_edge_torch_nightly-0.2.0.dev20240802.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-ai_edge_torch_nightly-0.2.0.dev20240802.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.2.0.dev20240802.dist-info/RECORD,,
+ai_edge_torch_nightly-0.2.0.dev20240803.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.2.0.dev20240803.dist-info/METADATA,sha256=3bNPO39qaqNIfjrr2bg74nLptaDTWtbEss9DTcYn5E0,1889
+ai_edge_torch_nightly-0.2.0.dev20240803.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+ai_edge_torch_nightly-0.2.0.dev20240803.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.2.0.dev20240803.dist-info/RECORD,,