ai-edge-torch-nightly 0.4.0.dev20250401__py3-none-any.whl → 0.4.0.dev20250404__py3-none-any.whl

ai_edge_torch/generative/test/test_custom_dus.py

@@ -15,6 +15,7 @@
 
 """A suite of tests to validate the Dynamic Update Slice Custom Op."""
 
+from ai_edge_torch.generative.custom_ops.dynamic_update_slice import dynamic_update_slice
 import torch
 from torch import nn
 

ai_edge_torch/odml_torch/lowerings/utils.py

@@ -14,10 +14,10 @@
 # ==============================================================================
 """Utilities for building MLIR lowerings."""
 
+from collections.abc import Callable
 import functools
 import numbers
-from typing import Any
-from typing import Optional
+from typing import Any, Optional, 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
@@ -222,3 +222,48 @@ def convert_int_to_float(t: ir.Value) -> ir.Value:
     return stablehlo.convert(
         ir.RankedTensorType.get(t.type.shape, ir.F64Type.get()), t
     )
+
+
+# IR Helpers
+IrValues = Union[ir.Value, tuple[ir.Value, ...]]
+
+
+# Non-canonicalized dtype to IR type mapping.
+_numpy_dtype_to_ir_type: dict[np.dtype, Callable[[], ir.Type]] = {
+    np.dtype(np.bool_): functools.partial(ir.IntegerType.get_signless, 1),
+    np.dtype(np.int8): functools.partial(ir.IntegerType.get_signless, 8),
+    np.dtype(np.int16): functools.partial(ir.IntegerType.get_signless, 16),
+    np.dtype(np.int32): functools.partial(ir.IntegerType.get_signless, 32),
+    np.dtype(np.int64): functools.partial(ir.IntegerType.get_signless, 64),
+    np.dtype(np.uint8): functools.partial(ir.IntegerType.get_unsigned, 8),
+    np.dtype(np.uint16): functools.partial(ir.IntegerType.get_unsigned, 16),
+    np.dtype(np.uint32): functools.partial(ir.IntegerType.get_unsigned, 32),
+    np.dtype(np.uint64): functools.partial(ir.IntegerType.get_unsigned, 64),
+    np.dtype(np.float16): ir.F16Type.get,
+    np.dtype(np.float32): ir.F32Type.get,
+    np.dtype(np.float64): ir.F64Type.get,
+    np.dtype(np.complex64): lambda: ir.ComplexType.get(ir.F32Type.get()),
+    np.dtype(np.complex128): lambda: ir.ComplexType.get(ir.F64Type.get()),
+}
+
+
+def numpy_dtype_to_ir_type(dtype: np.dtype | np.generic) -> ir.Type:
+  assert isinstance(dtype, (np.dtype, np.generic)), type(dtype)
+  dtype = np.dtype(dtype)
+  try:
+    ir_type_factory = _numpy_dtype_to_ir_type[dtype]
+  except KeyError as err:
+    raise TypeError(
+        f"No numpy_dtype_to_ir_type handler for dtype: {dtype}"
+    ) from err
+  return ir_type_factory()
+
+
+def numpy_array_constant(x: np.ndarray | np.generic) -> IrValues:
+  element_type = numpy_dtype_to_ir_type(x.dtype)
+  shape = x.shape
+  if x.dtype == np.bool_:
+    x = np.packbits(x, bitorder="little")  # type: ignore
+  x = np.ascontiguousarray(x)
+  attr = ir.DenseElementsAttr.get(x, type=element_type, shape=shape)  # type: ignore
+  return stablehlo.constant(attr)

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.4.0.dev20250401.dist-info → ai_edge_torch_nightly-0.4.0.dev20250404.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.4.0.dev20250401
+Version: 0.4.0.dev20250404
 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.dev20250401.dist-info → ai_edge_torch_nightly-0.4.0.dev20250404.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=alLtEtREkR7OYUxXcDYdLAFZ0Ba1BRJ3h2R9GVa6iBE,706
+ai_edge_torch/version.py,sha256=5Y5PWe5kvjSSYTDe16N_sj2O1E-n6hrbOjUlSz4mABA,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
@@ -176,7 +176,7 @@ ai_edge_torch/generative/quantize/quant_recipe_utils.py,sha256=4fgmP_GgeiFUOkIaC
 ai_edge_torch/generative/quantize/quant_recipes.py,sha256=0Kvr_o7pbMnE8VMe6Ml0FBxkHM6RJ3C14B2I1mjItjc,2030
 ai_edge_torch/generative/quantize/supported_schemes.py,sha256=FjdycEOvxRgBmQdZVufetPvkDoD7rUowIOSKV9oV5Kk,1418
 ai_edge_torch/generative/test/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
-ai_edge_torch/generative/test/test_custom_dus.py,sha256=ifgnUCWihT59eFdLrlc5_j9sWygEKclU6Iqw6zdlgeI,3177
+ai_edge_torch/generative/test/test_custom_dus.py,sha256=MjIhTvkTko872M35XMciobvICcDWTcIDJ3rociko-wM,3267
 ai_edge_torch/generative/test/test_kv_cache.py,sha256=MBPS-0bDXB0tQSKHa1XwDQeVIfabRbc8JQA99h9fzlQ,5961
 ai_edge_torch/generative/test/test_loader.py,sha256=9mQUeeZKOVApOWSWl2cN9c10axZjMKM1-0Zd823CCS4,3449
 ai_edge_torch/generative/test/test_lora.py,sha256=6QIM6RLTc2HrodGpp_aS3OxM9Rco2KAzEnYgotkg41M,5310
@@ -232,7 +232,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=tIyZiSy2Rbtea9ZlRYUfTaYE5vW_lAU6itT6_rUp8Qg,7028
+ai_edge_torch/odml_torch/lowerings/utils.py,sha256=-TzK1igPgR38oZkU1iPh-DZhlKVwuBtGWVC-y81PXzY,8935
 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
@@ -242,8 +242,8 @@ ai_edge_torch/testing/__init__.py,sha256=_yGgvnBZWb7T3IN3mc4x1sS4vM96HZwM8pwIcPG
 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.dev20250401.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.4.0.dev20250401.dist-info/METADATA,sha256=DiDVH9ci6trdmfoXY8vGF1P5g658z6tB9ijInpoCHlQ,1966
-ai_edge_torch_nightly-0.4.0.dev20250401.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.4.0.dev20250401.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.4.0.dev20250401.dist-info/RECORD,,
+ai_edge_torch_nightly-0.4.0.dev20250404.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.4.0.dev20250404.dist-info/METADATA,sha256=jpYvoXGbLlhrXvFTQ_mgRIp3skG8vmKU4966Je4VDNU,1966
+ai_edge_torch_nightly-0.4.0.dev20250404.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.4.0.dev20250404.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.4.0.dev20250404.dist-info/RECORD,,