ai-edge-torch-nightly 0.4.0.dev20250312__py3-none-any.whl → 0.4.0.dev20250314__py3-none-any.whl

ai_edge_torch/generative/examples/gemma/gemma2.py

@@ -247,6 +247,9 @@ def get_model_config_2b(kv_cache_max_len: int = 1024) -> cfg.ModelConfig:
         rotary_base=10000,
         rotary_percentage=1.0,
         qkv_transpose_before_split=True,
+        # The safetensors from HF is not using the interleaved qkv format, so
+        # we need to disable interleaving here in the model config.
+        qkv_fused_interleaved=False,
         logit_softcap=50.0,
         sliding_window_size=4096,
         attn_type=(

ai_edge_torch/odml_torch/lowerings/utils.py

@@ -18,14 +18,16 @@ import functools
 import numbers
 from typing import Any
 from typing import Optional
-
+from ai_edge_torch.odml_torch import export_utils
 from jax._src.lib.mlir import ir
 from jax._src.lib.mlir.dialects import hlo as stablehlo
 import numpy as np
 import torch
+import torch.utils._pytree as pytree
 
 
-def torch_dtype_to_ir_element_type(dtype):
+def torch_dtype_to_ir_element_type(dtype) -> ir.Type:
+  """Builds ir.Type from torch dtype."""
   ty_get = {
       torch.double: ir.F64Type.get,
       torch.float32: ir.F32Type.get,
@@ -39,6 +41,27 @@ def torch_dtype_to_ir_element_type(dtype):
   return ty_get()
 
 
+def node_meta_to_ir_types(node: torch.fx.Node) -> list[ir.Type]:
+  """Builds IR result types from torch FX node meta."""
+  tensor_meta = node.meta.get("tensor_meta") or node.meta.get("val")
+  if not tensor_meta:
+    raise RuntimeError(f"{node.name} does not have tensor meta")
+
+  tensor_meta_list, _ = pytree.tree_flatten(
+      [tensor_meta],
+      is_leaf=lambda x: hasattr(x, "dtype") and hasattr(x, "shape"),
+  )
+  results = []
+  for meta in tensor_meta_list:
+    shape = [
+        export_utils.IR_DYNAMIC if export_utils.is_torch_dynamic(dim) else dim
+        for dim in meta.shape
+    ]
+    elty = torch_dtype_to_ir_element_type(meta.dtype)
+    results.append(ir.RankedTensorType.get(shape, elty))
+  return results
+
+
 def splat(val, ty, shape=tuple(), *, loc: Optional[Any] = None):
   if isinstance(ty, ir.IntegerType):
     if ty.width == 1:

ai_edge_torch/testing/__init__.py

@@ -12,3 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
+from ai_edge_torch.testing import export
+from ai_edge_torch.testing import model_coverage
+
+export_with_tensor_inputs_only = export.export_with_tensor_inputs_only

ai_edge_torch/testing/export.py

@@ -0,0 +1,83 @@
+# Copyright 2025 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.
+# ==============================================================================
+"""Torch export utilities for testing."""
+
+from collections.abc import Callable
+from typing import Any
+
+import torch
+from torch.utils import _pytree as pytree
+
+
+def export_with_tensor_inputs_only(
+    model: Callable[..., Any],
+    args: tuple[Any, ...],
+    kwargs: dict[str, Any],
+) -> torch.export.ExportedProgram:
+  """Exports a PyTorch model, treating only tensor inputs as export inputs.
+
+  This function takes a PyTorch model and its input arguments (positional and
+  keyword) and exports it using `torch.export.export`. However, it modifies
+  the export process such that only the `torch.Tensor` arguments in the
+  inputs are considered as export inputs to the exported graph. All other
+  argument types (e.g., scalars, lists, tuples containing non-tensors) are
+  treated as constants.
+
+  This is useful for testing scenarios where you want to export a model but
+  want to avoid issues that might arise from non-tensor inputs
+  being treated as variables, or when you specifically want to focus on the
+  graph structure based on tensor operations.
+
+  Args:
+    model: The PyTorch `nn.Module` to be exported.
+    args: A tuple of positional arguments to be passed to the model's `forward`
+      method.
+    kwargs: A dictionary of keyword arguments to be passed to the model's
+      `forward` method.
+
+  Returns:
+    torch.export.ExportedProgram: The exported program representing the model
+    computation with only tensor inputs being export inputs.
+  """
+  flatten_args, treespec = pytree.tree_flatten([args, kwargs])
+
+  export_args = []
+  indices = []
+  for i, arg in enumerate(flatten_args):
+    if isinstance(arg, torch.Tensor):
+      export_args.append(arg)
+      indices.append(i)
+
+  class ModuleWrapper(torch.nn.Module):
+
+    def __init__(self, func, original_args, original_kwargs):
+      super().__init__()
+      self.original_args = list(flatten_args)
+      self.func = func
+
+    def forward(self, *export_args):
+      flatten_args = self.original_args.copy()
+      for i, arg in zip(indices, export_args):
+        flatten_args[i] = arg
+      args, kwargs = pytree.tree_unflatten(flatten_args, treespec)
+      return self.func(*args, **kwargs)
+
+  export_args = tuple(export_args)
+  export_kwargs = {}
+  return torch.export.export(
+      ModuleWrapper(model, args, kwargs).eval(),
+      export_args,
+      export_kwargs,
+  )

ai_edge_torch/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.4.0.dev20250312"
+__version__ = "0.4.0.dev20250314"

{ai_edge_torch_nightly-0.4.0.dev20250312.dist-info → ai_edge_torch_nightly-0.4.0.dev20250314.dist-info}/METADATA

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

@@ -2,7 +2,7 @@ ai_edge_torch/__init__.py,sha256=8sPR_5uXJA4NEE0nIwNdSl-ADOJEoR8hAgYvBQDY70Y,120
 ai_edge_torch/_config.py,sha256=AiqhbcheF7j_ozIGDLC89k1we95aVgFDa-tR6h7UI0s,2529
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/model.py,sha256=N-pNpTxzhaFGhWhnSGd70lBzb9VlEhTOq5mddU7bvvI,5542
-ai_edge_torch/version.py,sha256=PlBamGX8JQKmFS2RJl0lWF-mxslb0_eASGwwnezOHuY,706
+ai_edge_torch/version.py,sha256=PjlstuIJ-GlyKyFBMrwc7RQFRecNIkHpz5aIzvYNRKo,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=gpXQnifODU-mWxkUZw_3ov1lEYBw1SPVIcqj5k7pTGo,5550
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -57,7 +57,7 @@ ai_edge_torch/generative/examples/gemma/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIX
 ai_edge_torch/generative/examples/gemma/convert_gemma1_to_tflite.py,sha256=8HJi0cutxPstafVNs2LfBKdUzufVucje1Vrfjw_RS_g,2527
 ai_edge_torch/generative/examples/gemma/convert_gemma2_to_tflite.py,sha256=MX8fZhJJPZ5IoMiNHX0tLkRpHYqVuh4qhW0rkeIfmYw,2529
 ai_edge_torch/generative/examples/gemma/gemma1.py,sha256=w8oWYibZzvEvCDyp39EYyAWmjgJljhzdYPyFCfAWxZA,3497
-ai_edge_torch/generative/examples/gemma/gemma2.py,sha256=BOQ4zYKMXEX8Adly9-Yt6FBCtx1eW7IIppMSqmGvero,11356
+ai_edge_torch/generative/examples/gemma/gemma2.py,sha256=CK1lHw-YQPAr26KMdrYA6icQHvKH59yHAQ4eC4X636o,11539
 ai_edge_torch/generative/examples/gemma/verify_gemma1.py,sha256=ip-Gmk4CI5f0GWSdAIdrectxQWJ0t328KCsA4nfHuGg,1736
 ai_edge_torch/generative/examples/gemma/verify_gemma2.py,sha256=IoBhEMwH07-tFm5-U6F2hpCsI8xynglhq1x9tIOdaPQ,1322
 ai_edge_torch/generative/examples/gemma/verify_util.py,sha256=tR8RflXocDZqvuStyw9aFlzuiTllEC8rNnjrxms6_Is,5727
@@ -233,17 +233,18 @@ ai_edge_torch/odml_torch/lowerings/_quantized_decomposed.py,sha256=XDZ0zLej_XaQD
 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/registry.py,sha256=Tp2h11l5uTifO0aIkuUOWAF_ibEjmd65Xx99w3EXuGE,1924
-ai_edge_torch/odml_torch/lowerings/utils.py,sha256=pqM6mumpviFDHRaabp93CUAngzEZmWcAHl0nTDgyI2g,6167
+ai_edge_torch/odml_torch/lowerings/utils.py,sha256=tIyZiSy2Rbtea9ZlRYUfTaYE5vW_lAU6itT6_rUp8Qg,7028
 ai_edge_torch/odml_torch/passes/__init__.py,sha256=AVwIwUTMx7rXacKjGy4kwrtMd3XB2v_ncdc40KOjUqQ,1245
 ai_edge_torch/quantize/__init__.py,sha256=aB5dXot04bqyUhpsDFvxt9CIi15QAC4euvqOndJ0XLU,714
 ai_edge_torch/quantize/pt2e_quantizer.py,sha256=CKIEhs9jCcna64qj1jFH9zEbMbRdyeGV_TmSqEBPjes,15741
 ai_edge_torch/quantize/pt2e_quantizer_utils.py,sha256=nuO3w9gOj9sKcsTBBexVDw8UZnd06KsjNrFr_gyNaiA,36710
 ai_edge_torch/quantize/quant_config.py,sha256=U0KisSW-uZkoMJcy-ZP9W57p3tsa594fr9PphCRdO8o,3172
-ai_edge_torch/testing/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
+ai_edge_torch/testing/__init__.py,sha256=_yGgvnBZWb7T3IN3mc4x1sS4vM96HZwM8pwIcPGhPbo,833
+ai_edge_torch/testing/export.py,sha256=dguMa-aEi-WDPnmGBUs2IPdEmt2IVmHOELH19uiJ1uU,3014
 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.4.0.dev20250312.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.4.0.dev20250312.dist-info/METADATA,sha256=DTFxRv9pdU_Uy4wUIU8th5ZpgLzkYVFmM81SgjJAzAo,1966
-ai_edge_torch_nightly-0.4.0.dev20250312.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.4.0.dev20250312.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.4.0.dev20250312.dist-info/RECORD,,
+ai_edge_torch_nightly-0.4.0.dev20250314.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.4.0.dev20250314.dist-info/METADATA,sha256=n_c6T76WR-J-SCOmKKKzzuPoyM4i_2W2TO6ub8AuDw0,1966
+ai_edge_torch_nightly-0.4.0.dev20250314.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.4.0.dev20250314.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.4.0.dev20250314.dist-info/RECORD,,