ai-edge-torch-nightly 0.3.0.dev20241220__py3-none-any.whl → 0.3.0.dev20241221__py3-none-any.whl

ai_edge_torch/generative/examples/paligemma/decoder.py

@@ -19,6 +19,7 @@ from typing import Optional
 
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
 import ai_edge_torch.generative.layers.model_config as cfg
+import ai_edge_torch.generative.layers.rotary_position_embedding as rotary_pos_emb
 from ai_edge_torch.generative.utilities import model_builder
 import ai_edge_torch.generative.utilities.loader as loading_utils
 import torch
@@ -61,8 +62,12 @@ class Decoder(model_builder.DecoderOnlyModel):
     assert input_embeds is not None
 
     repo_pos = input_pos + 1  # PaliGemma position is 1-based.
-    cos, sin = self.rope_cache
-    rope = (cos.index_select(0, repo_pos), sin.index_select(0, repo_pos))
+    # ROPE parameters for all attn_configs are the same. Take the first one.
+    attn_config = self.config.block_config(0).attn_config
+    n_elem = int(attn_config.rotary_percentage * attn_config.head_dim)
+    rope = rotary_pos_emb.build_rope(
+        repo_pos, n_elem, attn_config.head_dim, attn_config.rotary_base
+    )
 
     # The first part of input_embeds are image embeddings. Diagonal causal mask
     # doesn't work here.

ai_edge_torch/generative/utilities/model_builder.py

@@ -107,8 +107,6 @@ class DecoderOnlyModel(nn.Module):
 
     # token embeddings of shape (b, t, n_embd)
     input_embeds = self.tok_embedding(tokens)
-    mask = self.mask_cache.index_select(2, input_pos)
-    mask = mask[:, :, :, : self.config.kv_cache_max]
 
     # ROPE parameters for all attn_configs are the same. Take the first one.
     attn_config = self.config.block_config(0).attn_config
@@ -117,6 +115,9 @@ class DecoderOnlyModel(nn.Module):
         input_pos, n_elem, attn_config.head_dim, attn_config.rotary_base
     )
 
+    mask = self.mask_cache.index_select(2, input_pos)
+    mask = mask[:, :, :, : self.config.kv_cache_max]
+
     return self.forward_with_embeds(
         input_embeds, rope, mask, input_pos, kv_cache, export_config
     )

ai_edge_torch/odml_torch/export.py

@@ -198,7 +198,12 @@ class MlirLowered:
     # build, which may not have the same StableHLO version as what used in
     # TFLite converter. Therefore we always serialize MLIR module in VHLO.
     # TODO(b/362798610) Build MLIR pybinding in ai-edge-torch release.
-    target_version = stablehlo.get_minimum_version()
+    if stablehlo.get_api_version() < 9:
+      target_version = stablehlo.get_minimum_version()
+    else:
+      target_version = stablehlo.get_version_from_compatibility_requirement(
+          stablehlo.StablehloCompatibilityRequirement.WEEK_4
+      )
     module_bytecode = xla_extension.mlir.serialize_portable_artifact(
         self.module_bytecode, target_version
     )

ai_edge_torch/odml_torch/jax_bridge/__init__.py

@@ -12,4 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from ai_edge_torch.odml_torch.jax_bridge._wrap import wrap
+from ai_edge_torch.odml_torch.jax_bridge import _wrap
+from ai_edge_torch.odml_torch.jax_bridge import utils
+
+wrap = _wrap.wrap

ai_edge_torch/odml_torch/lowerings/__init__.py

@@ -18,6 +18,7 @@ from . import _convolution
 from . import _jax_lowerings
 from . import _layer_norm
 from . import _quantized_decomposed
+from . import _rand
 from . import context
 from . import registry
 from . import utils

ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py

@@ -26,6 +26,7 @@ import torch_xla2.ops.ops_registry  # Import to load torch_xla2 ops
 
 LoweringContext = context.LoweringContext
 
+
 @functools.cache
 def _log_usage(op):
   logging.warning("Use jax lowering: %s", str(op))
@@ -184,8 +185,6 @@ lower_by_torch_xla2(torch.ops.aten.permute_copy)
 lower_by_torch_xla2(torch.ops.aten.pixel_shuffle)
 lower_by_torch_xla2(torch.ops.aten.pow)
 lower_by_torch_xla2(torch.ops.aten.prod)
-lower_by_torch_xla2(torch.ops.aten.rand)
-lower_by_torch_xla2(torch.ops.aten.randn)
 lower_by_torch_xla2(torch.ops.aten.reciprocal)
 lower_by_torch_xla2(torch.ops.aten.reflection_pad1d)
 lower_by_torch_xla2(torch.ops.aten.relu)

ai_edge_torch/odml_torch/lowerings/_rand.py

@@ -0,0 +1,142 @@
+# 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.
+# ==============================================================================
+import uuid
+
+from ai_edge_torch.odml_torch import export_utils
+from ai_edge_torch.odml_torch.lowerings import context
+from ai_edge_torch.odml_torch.lowerings import registry
+from jax._src.lib.mlir import ir
+from jax._src.lib.mlir.dialects import func
+from jax._src.lib.mlir.dialects import hlo as stablehlo
+import numpy as np
+import torch
+import torch.utils._pytree as pytree
+
+LoweringContext = context.LoweringContext
+lower = registry.lower
+
+
+def _random_lowering(
+    lctx: LoweringContext,
+    size: list[int],
+    generator,
+    dtype: torch.dtype,
+    rand_tensor,
+    composite_name: str,
+):
+  if dtype is None:
+    dtype = torch.float32
+
+  rand_tensor = rand_tensor.type(dtype)
+  data = rand_tensor.detach().numpy()
+
+  shape, _ = pytree.tree_flatten(size)
+  elty = export_utils.torch_dtype_to_ir_element_type(dtype)
+
+  decomp_name = f"{composite_name}.impl_{uuid.uuid4().hex[:8]}"
+
+  with ir.InsertionPoint(lctx.ir_module.body):
+
+    @func.FuncOp.from_py_func(
+        ir.RankedTensorType.get(
+            [len(shape)],
+            ir.IntegerType.get_signless(32),
+        ),
+        name=decomp_name,
+    )
+    def _rand_impl(_):
+      return [stablehlo.constant(ir.DenseElementsAttr.get(data))]
+
+  seed, seed2 = (
+      torch.randint(
+          torch.iinfo(torch.int64).min,
+          torch.iinfo(torch.int64).max,
+          (2,),
+          dtype=torch.int64,
+          generator=generator,
+      )
+      .detach()
+      .numpy()
+  )
+
+  shape_ = stablehlo.constant(
+      ir.DenseElementsAttr.get(np.array(shape, dtype=np.int32))
+  )
+  return stablehlo.CompositeOp(
+      result=[ir.RankedTensorType.get(shape, elty)],
+      inputs=[shape_],
+      name=composite_name,
+      composite_attributes=ir.DictAttr.get({
+          "seed": ir.IntegerAttr.get(ir.IntegerType.get_signless(64), seed),
+          "seed2": ir.IntegerAttr.get(ir.IntegerType.get_signless(64), seed2),
+      }),
+      decomposition=decomp_name,
+  ).results[0]
+
+
+# Schema:
+# - aten::rand(SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None,
+#     Device? device=None, bool? pin_memory=None) -> Tensor
+# - aten::rand.generator(SymInt[] size, *, Generator? generator,
+#     ScalarType? dtype=None, Layout? layout=None, Device? device=None,
+#     bool? pin_memory=None) -> Tensor
+@registry.lower(torch.ops.aten.rand)
+def _aten_rand(
+    lctx: LoweringContext,
+    size,
+    generator=None,
+    dtype=None,
+    layout=torch.strided,
+    device=None,
+    pin_memory=False,
+):
+  return _random_lowering(
+      lctx,
+      size,
+      generator,
+      dtype,
+      rand_tensor=torch.ops.aten.rand.generator(
+          size, generator=generator, dtype=dtype
+      ),
+      composite_name="odml.random_uniform",
+  )
+
+
+# Schema:
+# - aten::randn(SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None,
+#     Device? device=None, bool? pin_memory=None) -> Tensor
+# - aten::randn.generator(SymInt[] size, *, Generator? generator,
+#     ScalarType? dtype=None, Layout? layout=None, Device? device=None,
+#     bool? pin_memory=None) -> Tensor
+@registry.lower(torch.ops.aten.randn)
+def _aten_randn(
+    lctx: LoweringContext,
+    size,
+    generator=None,
+    dtype=None,
+    layout=torch.strided,
+    device=None,
+    pin_memory=False,
+):
+  return _random_lowering(
+      lctx,
+      size,
+      generator,
+      dtype,
+      rand_tensor=torch.ops.aten.randn.generator(
+          size, generator=generator, dtype=dtype
+      ),
+      composite_name="odml.random_standard_normal",
+  )

ai_edge_torch/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.3.0.dev20241220"
+__version__ = "0.3.0.dev20241221"

{ai_edge_torch_nightly-0.3.0.dev20241220.dist-info → ai_edge_torch_nightly-0.3.0.dev20241221.dist-info}/METADATA

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

@@ -3,7 +3,7 @@ ai_edge_torch/_config.py,sha256=QIrerb6uHMahRvMilmhodJ_6jfiRps3qgLOBeidPnS4,1614
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/fx_pass_base.py,sha256=518ziQ0TUxqum2qZXqlD8qr65pHPh8ZNLnwFC6zvK3k,4253
 ai_edge_torch/model.py,sha256=N-pNpTxzhaFGhWhnSGd70lBzb9VlEhTOq5mddU7bvvI,5542
-ai_edge_torch/version.py,sha256=xD-MWAEa1ROHhyF3rY7MaL28xsuON0aJwaiXbJ04qfc,706
+ai_edge_torch/version.py,sha256=4pSrONNJgkt6DeTfleRz5DpcHts3SW-iInT2ibr1t9A,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=SzbR16V2JEfkCjjPwRVAFUbFnzu-_1iHPKgGT9Yz7gQ,5678
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -64,7 +64,7 @@ ai_edge_torch/generative/examples/openelm/openelm.py,sha256=sIJ8Ie1oxFrJM-1jvv2u
 ai_edge_torch/generative/examples/openelm/verify.py,sha256=VkigoqhAr8ew95neb3TifYv-SLOSheaWKv2AH0iKDrc,2441
 ai_edge_torch/generative/examples/paligemma/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/examples/paligemma/convert_to_tflite.py,sha256=rPFqcsv8RHvjmgfBW9OL6EKxMtVX-ySjBsMP4N8FErk,2816
-ai_edge_torch/generative/examples/paligemma/decoder.py,sha256=eICKQkJsJuEUkuvn5ymUsI9CGB-oNbgV7VH7BlmklfQ,4961
+ai_edge_torch/generative/examples/paligemma/decoder.py,sha256=HDDTd4F0kOurhXyqikP5umdY0gVm-FHA1ysaKcz88CM,5261
 ai_edge_torch/generative/examples/paligemma/image_encoder.py,sha256=yKPWG8aBp-GuzeyQntlzwTTcGBBjvUywVGRjnlNprmo,5574
 ai_edge_torch/generative/examples/paligemma/paligemma.py,sha256=nDyI-wUFJSawu57uLbFENei5l4cciqZ8lM5S5beN0FU,5604
 ai_edge_torch/generative/examples/paligemma/verify.py,sha256=Bkbgy-GFjnMNYjduWUM7YLWarPTwmj1v38eHY-PdBlM,4874
@@ -147,7 +147,7 @@ ai_edge_torch/generative/utilities/__init__.py,sha256=-_jxnnFnCgnTU4oTm4MnRsvL5l
 ai_edge_torch/generative/utilities/converter.py,sha256=hIwWUWjgPvWLATtsYYG6RWbFQWhOr2RpPlMrd-4Am9U,5959
 ai_edge_torch/generative/utilities/dynamic_update_slice.py,sha256=e2mhx-Vp8sUK4EXoPtpZLSx3TViqLAKs67EhKcXBjAQ,2121
 ai_edge_torch/generative/utilities/loader.py,sha256=A3SOjPXp--AsvoP1hqj5QKWE4sgxoFc3H5EBUz_Eogc,13531
-ai_edge_torch/generative/utilities/model_builder.py,sha256=q82-1E2zYlzpbFW6Vw-MWrJivRXHKpRh8jUxpR-w0sY,6349
+ai_edge_torch/generative/utilities/model_builder.py,sha256=plKHp5csnZpx3GQ1SYTqFpdoaxTVcwXgCmzO5N6ya6I,6350
 ai_edge_torch/generative/utilities/moonshine_loader.py,sha256=_RpFabSqtGH5PHiP3_1f6QfO14qMADUxr_HGRlVDFB0,4891
 ai_edge_torch/generative/utilities/stable_diffusion_loader.py,sha256=dqPD9qRXEWtU3ombslOC-BE2l_dMwHoCNu7NsIJhsso,36158
 ai_edge_torch/generative/utilities/t5_loader.py,sha256=tEsfy8-ymzbbjOIc-oesXF3yGyyWtJgFXn2s7VOavt8,16961
@@ -169,7 +169,7 @@ ai_edge_torch/lowertools/translate_recipe.py,sha256=ymkBpFqAUiupRWqrPOWiVphKcXR1
 ai_edge_torch/odml_torch/__init__.py,sha256=S8jOzE9nLof-6es3XDiGJRN-9H_XTxsVm9dE7lD3RWo,812
 ai_edge_torch/odml_torch/_torch_future.py,sha256=AJ0klpsbu2ZBTfiZlqSOoaYzBVITt40a1fYN8xKkEPw,3044
 ai_edge_torch/odml_torch/_torch_library.py,sha256=Lw1gqL2HWNRspdTwNhIkYAHDyafHedHtkXyKKxn-Wss,805
-ai_edge_torch/odml_torch/export.py,sha256=Wc_JM7U2IjZeBmXA6t1AZxREGOWjZ6EB-PIhEevWWeU,13207
+ai_edge_torch/odml_torch/export.py,sha256=QzOPmcNPB7R-KhhPEP0oGVbDRgGPptIxRSoz3S8py9I,13405
 ai_edge_torch/odml_torch/export_utils.py,sha256=QeA37Irlty6AiIBuqmHmJgn3lqahBQ5xsh6IKRoKm1g,4774
 ai_edge_torch/odml_torch/tf_integration.py,sha256=lTFJPPEijLPFmn6qq2jbpVTQOo0YaOTK36kK6rCiyIE,5956
 ai_edge_torch/odml_torch/composite/__init__.py,sha256=71GM_gDZxJyo38ZSoYSwhZX3xKA9rknO93JS9kw9w_c,778
@@ -178,16 +178,17 @@ ai_edge_torch/odml_torch/composite/stablehlo_composite_builder.py,sha256=2Y52E_g
 ai_edge_torch/odml_torch/debuginfo/__init__.py,sha256=9ag6-WWRG50rPCtIV7OpIokEKu2YRyGlMZZqVPWUH6g,762
 ai_edge_torch/odml_torch/debuginfo/_build.py,sha256=1xCXOs3-9UcsOyLFH0uyQwLu7c06iYFTo0NQ7Ckbl2I,1465
 ai_edge_torch/odml_torch/debuginfo/_op_polyfill.py,sha256=IvOBQyROI9WHS3umHRxsDW-1YElU9BPWzKtJA2eKWOI,1739
-ai_edge_torch/odml_torch/jax_bridge/__init__.py,sha256=Jco5zvejxuyl9xHQxZICAKbkgH7x38qPlwUUpD7S15Q,730
+ai_edge_torch/odml_torch/jax_bridge/__init__.py,sha256=e9Oa4J3An9FYr3zM0OzjzyNNitEeg-IoBUGNfUxsDSA,798
 ai_edge_torch/odml_torch/jax_bridge/_wrap.py,sha256=LqwZ1vCJTSOzgzvH8LUAN-sAkF-l_pGj1AMEIzAqHCA,6638
 ai_edge_torch/odml_torch/jax_bridge/utils.py,sha256=T8isGc896VrHZ6c_L5pYmLpolQ7ibcOlgWfPuVFPzIg,2264
-ai_edge_torch/odml_torch/lowerings/__init__.py,sha256=0GytV1dGnqe1mKityqQDNFNS8T4QBg3UZuRJcGHwGyA,993
+ai_edge_torch/odml_torch/lowerings/__init__.py,sha256=GWFl7WWgExLXu6FEYxnig5_g6hd_Sfnl8690uFg2-CU,1013
 ai_edge_torch/odml_torch/lowerings/_basic.py,sha256=8mZTp_ybcMO3tDRQdlDP68BVeTw560XsTR4XH-ldTdc,9987
 ai_edge_torch/odml_torch/lowerings/_batch_norm.py,sha256=PaLI0BB6pdBW1VyfW8VTOT_Be-ZcqYdNOsyfzKfq8Cg,2064
 ai_edge_torch/odml_torch/lowerings/_convolution.py,sha256=Q0aDzyUcZMoSzSbOU-r3LJMgPe6fble0QwdYVIOHHHk,6887
-ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=OVmlPGwyhDXKhmG4SAeEsa6iLpJHEHV_jKqwfjYvetA,11643
+ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=CJHWkmY4aAVQ5dmFsVc3Ox9TPkoLSNOfa96psD4CLRo,11561
 ai_edge_torch/odml_torch/lowerings/_layer_norm.py,sha256=khJIvDVk2s332Nd2Be-5dM6-wp5DGff61HCV5lskHmQ,3011
 ai_edge_torch/odml_torch/lowerings/_quantized_decomposed.py,sha256=XDZ0zLej_XaQDJnaAAxhNFAd7NfQm5SOVEp_nno_krA,6178
+ai_edge_torch/odml_torch/lowerings/_rand.py,sha256=g6SuqDkuC6hD35lyP1-5H7ASDIzPSmKukeNT5naZSv8,4133
 ai_edge_torch/odml_torch/lowerings/context.py,sha256=jslcCv7r_HtImSRTxJwHAUV_QCu9Jub51lovmoBkmFA,1295
 ai_edge_torch/odml_torch/lowerings/decomp.py,sha256=UoJeZVcr4zAN_11i-HzfOhxGCxUm-7b1JXPVBxR2hSs,2414
 ai_edge_torch/odml_torch/lowerings/registry.py,sha256=Tp2h11l5uTifO0aIkuUOWAF_ibEjmd65Xx99w3EXuGE,1924
@@ -200,8 +201,8 @@ ai_edge_torch/quantize/quant_config.py,sha256=U0KisSW-uZkoMJcy-ZP9W57p3tsa594fr9
 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=UPB448aMDUyC0HNYVqio2rcJPnDN0tBQMP08J6vPYew,4718
-ai_edge_torch_nightly-0.3.0.dev20241220.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20241220.dist-info/METADATA,sha256=PfyYhqbf7VEibw2TEDRb8tBOIPG9dfXhT9tNNou_iZg,1966
-ai_edge_torch_nightly-0.3.0.dev20241220.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.3.0.dev20241220.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20241220.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20241221.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20241221.dist-info/METADATA,sha256=_mQiElLiIpig6KWylK15amdyQP57haDyWH4Xaqqt_Ls,1966
+ai_edge_torch_nightly-0.3.0.dev20241221.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.3.0.dev20241221.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20241221.dist-info/RECORD,,