ai-edge-torch-nightly 0.3.0.dev20241211__py3-none-any.whl → 0.3.0.dev20241212__py3-none-any.whl

ai_edge_torch/generative/examples/paligemma/decoder.py

@@ -15,6 +15,8 @@
 
 """Example of building a decoder of PaliGemma 3B model which is Gemma1."""
 
+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
 from ai_edge_torch.generative.utilities import model_builder
@@ -51,6 +53,7 @@ class Decoder(model_builder.DecoderOnlyModel):
       input_pos: torch.Tensor,
       kv_cache: kv_utils.KVCache,
       input_embeds: torch.Tensor = None,
+      export_config: Optional[model_builder.ExportConfig] = None,
   ) -> dict[torch.Tensor, kv_utils.KVCache]:
     if input_embeds is None:
       return super().forward(tokens, input_pos, kv_cache)

ai_edge_torch/generative/examples/paligemma/paligemma.py

@@ -16,11 +16,13 @@
 """Example of building a full-stack of PaliGemma model."""
 
 from dataclasses import dataclass
+from typing import Optional
 
 from ai_edge_torch.generative.examples.paligemma import decoder
 from ai_edge_torch.generative.examples.paligemma import image_encoder
 import ai_edge_torch.generative.layers.kv_cache as kv_utils
 import ai_edge_torch.generative.layers.model_config as cfg
+from ai_edge_torch.generative.utilities import model_builder
 import ai_edge_torch.generative.utilities.loader as loading_utils
 import torch
 from torch import nn
@@ -67,9 +69,16 @@ class PaliGemma(nn.Module):
       input_pos: torch.Tensor,
       kv_cache: kv_utils.KVCache,
       pixel_values: torch.Tensor = None,
+      export_config: Optional[model_builder.ExportConfig] = None,
   ) -> dict[torch.Tensor, kv_utils.KVCache]:
     if pixel_values is None:
-      return self.decoder(tokens, input_pos, kv_cache)
+      return self.decoder(
+          tokens=tokens,
+          input_pos=input_pos,
+          kv_cache=kv_cache,
+          input_embeds=None,
+          export_config=export_config
+      )
 
     input_embeds = self.decoder.tok_embedding(tokens)
 
@@ -100,6 +109,7 @@ class PaliGemma(nn.Module):
         input_pos=input_pos,
         kv_cache=kv_cache,
         input_embeds=input_embeds,
+        export_config=export_config,
     )
 
 

ai_edge_torch/generative/layers/normalization.py

@@ -190,14 +190,12 @@ def group_norm_with_hlfb(
   """
   x = torch.permute(x, (0, 2, 3, 1))
 
-  # TODO: b/366544750 - Change "reduction_axes" field as an array, rather than
-  # int32 when the bug is fixed.
   builder = StableHLOCompositeBuilder(
       name="odml.group_norm",
       attr={
           "num_groups": num_groups,
           "epsilon": eps,
-          "reduction_axes": 3,
+          "reduction_axes": [3],
           "channel_axis": 3,
       },
   )

ai_edge_torch/generative/utilities/transformers_verifier.py

@@ -29,7 +29,7 @@ class TransformersModelWrapper(verifier.ModelWrapper):
   an object with `logits` field.
 
   Transformers models get `max_new_tokens` settings for generate() via
-  ExportConfig.
+  GenerationConfig.
   """
 
   def forward(self, tokens: torch.Tensor) -> torch.Tensor:
@@ -38,5 +38,5 @@ class TransformersModelWrapper(verifier.ModelWrapper):
   def generate(
       self, inputs: torch.Tensor, max_new_tokens: int
   ) -> torch.IntTensor:
-    export_config = transformers.ExportConfig(max_new_tokens=max_new_tokens)
-    return self.model.generate(inputs=inputs, generation_config=export_config)
+    gen_config = transformers.GenerationConfig(max_new_tokens=max_new_tokens)
+    return self.model.generate(inputs=inputs, generation_config=gen_config)

ai_edge_torch/generative/utilities/verifier.py

@@ -115,7 +115,7 @@ class ReauthoredModelWrapper(ModelWrapper):
     # pixel_values only when it is not None. Otherwise, it may raise an error.
     if pixel_values is None:
       output = self.model.forward(
-          tokens, input_pos, kv_cache, self.export_config
+          tokens, input_pos, kv_cache, export_config=self.export_config
       )
     else:
       output = self.model.forward(

ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py

@@ -16,12 +16,15 @@ import functools
 import logging
 
 from ai_edge_torch.odml_torch import jax_bridge
+from ai_edge_torch.odml_torch.lowerings import context
+from ai_edge_torch.odml_torch.lowerings import registry
+import jax.numpy as jnp
+from jax._src.lib.mlir import ir
 import torch
 import torch_xla2.ops.jaten  # Import to load torch_xla2 ops
 import torch_xla2.ops.ops_registry  # Import to load torch_xla2 ops
 
-from . import registry
-
+LoweringContext = context.LoweringContext
 
 @functools.cache
 def _log_usage(op):
@@ -258,3 +261,26 @@ def _aten_copy(self, *args, **kwargs):
 @lower_by_jax(torch.ops.aten.copy, ir_input_names=["src"])
 def _aten_copy(self, src, **kwargs):
   return _TORCH_XLA2_IMPLS[torch.ops.aten.copy](self, src)
+
+
+# Schema:
+#   - aten::einsum(str equation, Tensor[] tensors, *, int[]? path=None)
+#       -> Tensor
+# Torch Reference:
+#   - https://pytorch.org/docs/stable/generated/torch.einsum.html
+#   - https://github.com/pytorch/pytorch/blob/1b3f8b75896720e88362cbec7db32abc52afa83e/aten/src/ATen/native/Linear.cpp#L255
+@registry.lower(torch.ops.aten.einsum.default)
+def _aten_einsum_default(
+    lctx: LoweringContext,
+    equation: str,
+    tensors: list[ir.Value],
+    path=None,
+):
+  _log_usage(torch.ops.aten.einsum.default)
+
+  @jax_bridge.wrap
+  def jax_lowering(operands):
+    # Ignore the input path and let JAX determine the path.
+    return jnp.einsum(equation, *operands, optimize="optimal")
+
+  return jax_lowering(lctx, tuple(tensors))

ai_edge_torch/odml_torch/lowerings/decomp.py

@@ -46,7 +46,13 @@ def decompositions():
 
   torch._decomp.remove_decompositions(
       decompositions,
-      [torch.ops.aten.roll],
+      [
+          torch.ops.aten.roll,
+          # Torch's default einsum impl/decompositions is less efficient and
+          # optimized through converter than JAX's impl. Disable einsum
+          # decomposition to use JAX bridge for a more efficient lowering.
+          torch.ops.aten.einsum.default,
+      ],
   )
 
   # Override _safe_softmax decompositions with regular softmax.

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.3.0.dev20241211.dist-info → ai_edge_torch_nightly-0.3.0.dev20241212.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.3.0.dev20241211
+Version: 0.3.0.dev20241212
 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
@@ -29,7 +29,7 @@ Requires-Dist: safetensors
 Requires-Dist: tabulate
 Requires-Dist: torch>=2.4.0
 Requires-Dist: torch-xla>=2.4.0
-Requires-Dist: tf-nightly>=2.19.0.dev20241121
+Requires-Dist: tf-nightly>=2.19.0.dev20241201
 Requires-Dist: ai-edge-litert-nightly
 Requires-Dist: ai-edge-quantizer-nightly
 

{ai_edge_torch_nightly-0.3.0.dev20241211.dist-info → ai_edge_torch_nightly-0.3.0.dev20241212.dist-info}/RECORD

@@ -3,7 +3,7 @@ ai_edge_torch/config.py,sha256=FMWeCH2b7HYILBvaI1iZNnYCO4WAhDOwBZBmIE-xrF0,909
 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=_uS2Df0H-aUbz-7M-gLxfjDVOJxr03EeNDfbVC_cBrE,706
+ai_edge_torch/version.py,sha256=QnJ2_alMOUe5ea0vTpY7AIBr8eoHvuwKaaj917g5DFA,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=HwzfRx_DX5TLtPqwEH1_NOm38_INvHzHl4_mX67KOdQ,5448
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -64,9 +64,9 @@ 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=f_A3GWcLrP0nRq2Tq-fThfXIQVJ-EYWoExYLO_6iVIQ,4866
+ai_edge_torch/generative/examples/paligemma/decoder.py,sha256=eICKQkJsJuEUkuvn5ymUsI9CGB-oNbgV7VH7BlmklfQ,4961
 ai_edge_torch/generative/examples/paligemma/image_encoder.py,sha256=yKPWG8aBp-GuzeyQntlzwTTcGBBjvUywVGRjnlNprmo,5574
-ai_edge_torch/generative/examples/paligemma/paligemma.py,sha256=pIjsS-IUFevRjFA9153YT1vtWXATGWHsgVQQX_nWaZQ,5280
+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
 ai_edge_torch/generative/examples/paligemma/verify_decoder.py,sha256=al5wMPWri4IRVWrLmCplPi6uoCzwh0vBHMGnCt-XUqo,2690
 ai_edge_torch/generative/examples/paligemma/verify_image_encoder.py,sha256=pSekf1BybhieQz3cQx_llbRQHxczXbTqool8fOyGj_0,3114
@@ -121,7 +121,7 @@ ai_edge_torch/generative/layers/builder.py,sha256=LXGuSHIx6QZAzLFm7aJvlzoMPgQwbX
 ai_edge_torch/generative/layers/feed_forward.py,sha256=hdICat-8gW7-vxDAevJQ8NQ-mynllPiqLdXQMF6JMnc,4189
 ai_edge_torch/generative/layers/kv_cache.py,sha256=lbm-yJ1jGPtcgWS4C3FmSnB1IlxqDE7g0BLRh3PN4N4,6324
 ai_edge_torch/generative/layers/model_config.py,sha256=viX51T_naJ9sPpPxPoMnSueBPYE2zxWNOD0xn0f-_bM,7510
-ai_edge_torch/generative/layers/normalization.py,sha256=_2hps2m2MXEHQWbM-1B4he90hbq8wqOnIDIf-qXHhpc,7589
+ai_edge_torch/generative/layers/normalization.py,sha256=h2btgRHMMjOcyLm8adEmcT0pG6imq4QcWblKJK5MYXA,7479
 ai_edge_torch/generative/layers/rotary_position_embedding.py,sha256=xxWtlVsGGJkEyXC6PwznubyhJnLPEfSpHOORE_hgxss,2670
 ai_edge_torch/generative/layers/scaled_dot_product_attention.py,sha256=gXxh3papKy4FBpGEX7VyZ7rZ1Js6aHK70Q6DKrVSckY,4154
 ai_edge_torch/generative/layers/unet/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
@@ -151,8 +151,8 @@ ai_edge_torch/generative/utilities/model_builder.py,sha256=rfD6INxunvDVdiUfTUxD7
 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
-ai_edge_torch/generative/utilities/transformers_verifier.py,sha256=nHmI27ybu7lj8Ufw2LzmCwRDqEwNppIFNTx5ltLHIgE,1547
-ai_edge_torch/generative/utilities/verifier.py,sha256=1NcmT_55Sb5e5spnHab4x5wqJZi2CKKVtXuXgK3lE6Q,11927
+ai_edge_torch/generative/utilities/transformers_verifier.py,sha256=8sp9m_FMcXn7nqOrochtu2jIANkJKhnhIBUmH0ZTDR4,1549
+ai_edge_torch/generative/utilities/verifier.py,sha256=ESSA8W1EYNsd4ntwmXbr-dn-BcIS27hf53XL5RTwjEU,11941
 ai_edge_torch/hlfb/__init__.py,sha256=sH4um75na-O8tzxN6chFyp6Y4xnexsE7kUQpZySv6dE,735
 ai_edge_torch/hlfb/mark_pattern/__init__.py,sha256=cjTprggj_cuktSCm7-A25e7Shop3k63ylp7sdZmtZ8o,4790
 ai_edge_torch/hlfb/mark_pattern/passes.py,sha256=pjkKcI1nHECPluAt87cFBrt1DP0f3ge7rHq1NhCkBIE,1936
@@ -186,11 +186,11 @@ ai_edge_torch/odml_torch/lowerings/__init__.py,sha256=0GytV1dGnqe1mKityqQDNFNS8T
 ai_edge_torch/odml_torch/lowerings/_basic.py,sha256=ufvnaAh6rM_yfoc8ybI3VErHEVBv5W_p4iOe9slfwKM,9948
 ai_edge_torch/odml_torch/lowerings/_batch_norm.py,sha256=PaLI0BB6pdBW1VyfW8VTOT_Be-ZcqYdNOsyfzKfq8Cg,2064
 ai_edge_torch/odml_torch/lowerings/_convolution.py,sha256=v1VdKmL8YLJv3PR9VgyNghO83A25PpTzY2ZUAJqlq3Q,6847
-ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=4UyNyaR2W-vCOvj-P5lywQ1_RfLIxVE7J_GONI6CQvI,10718
+ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=OVmlPGwyhDXKhmG4SAeEsa6iLpJHEHV_jKqwfjYvetA,11643
 ai_edge_torch/odml_torch/lowerings/_layer_norm.py,sha256=1ePJs7oIdUkVdMddFsXMc53qTkEKqGz0ZhQQoNzBa10,2862
 ai_edge_torch/odml_torch/lowerings/_quantized_decomposed.py,sha256=rFmzqcdjYrwhcxH8j9zCFStPy21HFF7hkUV_GQ8FPAk,6056
 ai_edge_torch/odml_torch/lowerings/context.py,sha256=jslcCv7r_HtImSRTxJwHAUV_QCu9Jub51lovmoBkmFA,1295
-ai_edge_torch/odml_torch/lowerings/decomp.py,sha256=aR6JPFP2Iq-aR0qPxJEHehmAVTjiGhgQEoycZV_1vPY,2130
+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
 ai_edge_torch/odml_torch/lowerings/utils.py,sha256=pqM6mumpviFDHRaabp93CUAngzEZmWcAHl0nTDgyI2g,6167
 ai_edge_torch/odml_torch/passes/__init__.py,sha256=AVwIwUTMx7rXacKjGy4kwrtMd3XB2v_ncdc40KOjUqQ,1245
@@ -201,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.dev20241211.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20241211.dist-info/METADATA,sha256=Lyub5vadYf6Yu6mGY7l1PFk8Jg2rB36ojIBHm9CxhBM,1897
-ai_edge_torch_nightly-0.3.0.dev20241211.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.3.0.dev20241211.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20241211.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20241212.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20241212.dist-info/METADATA,sha256=tkJXZvoB1p4WKAKgK9Ql071JxwI7BwU3gKmdJR5jcrs,1897
+ai_edge_torch_nightly-0.3.0.dev20241212.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.3.0.dev20241212.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20241212.dist-info/RECORD,,