ai-edge-torch-nightly 0.6.0.dev20250606__py3-none-any.whl → 0.6.0.dev20250608__py3-none-any.whl

ai_edge_torch/odml_torch/lowerings/utils.py

@@ -17,7 +17,7 @@
 from collections.abc import Callable
 import functools
 import numbers
-from typing import Any, Optional, Union
+from typing import Any, Optional, Sequence, Union
 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
@@ -281,3 +281,33 @@ def convert_to_ir_value(
   if isinstance(value, ir.Value):
     return value
   raise TypeError(f"Unsupported type for conversion to ir.Value: {type(value)}")
+
+
+def convert_shape_to_ir_value(
+    shape: Sequence[int],
+) -> ir.Value:
+  # Check if all elements in the shape sequence are integers.
+  if not shape or all(isinstance(dim, int) for dim in shape):
+    # If all are integers, create a constant numpy array.
+    # Assuming int32 is the required type for TFLite shape tensors.
+    shape_ir_value = numpy_array_constant(np.array(shape, dtype=np.int32))
+  else:
+    # Handle mixed int and ir.Value shape sequence
+    processed_dims = []
+    for dim in shape:
+      if isinstance(dim, int):
+        # Convert int to a constant 1D tensor
+        shape_ir_value = numpy_array_constant(np.array([dim], dtype=np.int32))
+        processed_dims.append(shape_ir_value)
+      else:
+        assert isinstance(dim, ir.Value)
+        # Convert ir.Value to a constant 1D tensor
+        new_type = ir.RankedTensorType.get([1], dim.type.element_type)
+        reshape_dim = stablehlo.reshape(new_type, dim)
+        processed_dims.append(reshape_dim)
+
+    shape_ir_value = stablehlo.concatenate(
+        processed_dims,
+        dimension=0,
+    )
+  return shape_ir_value

ai_edge_torch/version.py

@@ -15,4 +15,4 @@
 
 # The next version of ai-edge-torch.
 # The minor version code should be bumped after every release.
-__version__ = "0.6.0.dev20250606"
+__version__ = "0.6.0.dev20250608"

{ai_edge_torch_nightly-0.6.0.dev20250606.dist-info → ai_edge_torch_nightly-0.6.0.dev20250608.dist-info}/METADATA

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

@@ -2,7 +2,7 @@ ai_edge_torch/__init__.py,sha256=lemyLCNoGYRnJsmDuGZu7qOqLbLqG6CGDFtu3ue1syU,129
 ai_edge_torch/_config.py,sha256=AiqhbcheF7j_ozIGDLC89k1we95aVgFDa-tR6h7UI0s,2529
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/model.py,sha256=A7loFu8jE9CsXsfMmHYZ-KDFJiaD8Kkqwm_9d3IVzk0,5638
-ai_edge_torch/version.py,sha256=w38CdzAqRuFGjEjznLvOmY9z1dEKo3ySTSdgY6aBdSU,806
+ai_edge_torch/version.py,sha256=BXtcfzSJXS87SYAsACZbMKtP4z7lF-KawEz0OLQFMSk,806
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=iQk3R-pLq4c1nfLqPB4xTRj78gghxPGzJCJtILLdg5o,6123
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -258,7 +258,7 @@ 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=wHIhaKGWxf_x4W750_nzKoYMBDlq2Fd6b05-XOEVyVQ,9465
+ai_edge_torch/odml_torch/lowerings/utils.py,sha256=YagY2IDqtJTpPxjPz2dhb3eyCFTpTSu3ptEPSvEuDtk,10574
 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
@@ -268,8 +268,8 @@ ai_edge_torch/testing/__init__.py,sha256=_yGgvnBZWb7T3IN3mc4x1sS4vM96HZwM8pwIcPG
 ai_edge_torch/testing/export.py,sha256=k5mGDGzwc23Z4zaIVDs8CNh-oOt64gsf9MS9NjhbPy4,3293
 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.6.0.dev20250606.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.6.0.dev20250606.dist-info/METADATA,sha256=dlZhk9uhgYBPi7Z3Mywcu34k4qnslMz6QHqqaaVA9_M,2074
-ai_edge_torch_nightly-0.6.0.dev20250606.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.6.0.dev20250606.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.6.0.dev20250606.dist-info/RECORD,,
+ai_edge_torch_nightly-0.6.0.dev20250608.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.6.0.dev20250608.dist-info/METADATA,sha256=2N-NXQ9DxVdBXhhR1Emb84854aBitqs6rTGp4DTlE-o,2074
+ai_edge_torch_nightly-0.6.0.dev20250608.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.6.0.dev20250608.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.6.0.dev20250608.dist-info/RECORD,,