ai-edge-torch-nightly 0.5.0.dev20250410__py3-none-any.whl → 0.5.0.dev20250412__py3-none-any.whl

ai_edge_torch/generative/layers/experimental/scaled_dot_product_attention.py

@@ -20,7 +20,7 @@ from typing import Optional
 
 from ai_edge_torch.generative.custom_ops import bmm_4d as bmm_lib
 from ai_edge_torch.generative.layers import kv_cache as kv_utils
-from ai_edge_torch.generative.layers.experimental import types
+from ai_edge_torch.generative.utilities import types
 from ai_edge_torch.hlfb import StableHLOCompositeBuilder
 from multipledispatch import dispatch
 import torch

ai_edge_torch/generative/layers/kv_cache.py

@@ -20,7 +20,7 @@ from typing import Any, List, Tuple
 
 from ai_edge_torch.generative.custom_ops.dynamic_update_slice import dynamic_update_slice
 from ai_edge_torch.generative.layers import model_config
-from ai_edge_torch.generative.layers.experimental import types
+from ai_edge_torch.generative.utilities import types
 import torch
 import torch.utils._pytree as pytree
 

ai_edge_torch/generative/layers/sdpa_with_kv_update.py

@@ -19,8 +19,8 @@ from ai_edge_torch.generative.layers import kv_cache as kv_utils
 from ai_edge_torch.generative.layers import scaled_dot_product_attention as sdpa_default
 from ai_edge_torch.generative.layers.experimental import kv_cache as kv_utils_experimental
 from ai_edge_torch.generative.layers.experimental import scaled_dot_product_attention as sdpa
-from ai_edge_torch.generative.layers.experimental import types
 import ai_edge_torch.generative.layers.model_config as cfg
+from ai_edge_torch.generative.utilities import types
 from multipledispatch import dispatch
 import torch
 

ai_edge_torch/version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.5.0.dev20250410"
+__version__ = "0.5.0.dev20250412"

{ai_edge_torch_nightly-0.5.0.dev20250410.dist-info → ai_edge_torch_nightly-0.5.0.dev20250412.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.5.0.dev20250410
+Version: 0.5.0.dev20250412
 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.5.0.dev20250410.dist-info → ai_edge_torch_nightly-0.5.0.dev20250412.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=dQvyVQmvNYF8n8HwlkY-9fdSo-n3_bdLO9EAZpJnC8s,706
+ai_edge_torch/version.py,sha256=mEMXseCOEc1eT2LhQ6YlU7_bA_5FIvGljV7Y1TT49h4,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=GPDsXhfECjDzOut4vh_d9qWcyfpxobFMBTsC7MyJbM0,5557
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -154,18 +154,17 @@ ai_edge_torch/generative/layers/attention.py,sha256=wLZ1jgUlcODBWgK3hnnhclHuuQDq
 ai_edge_torch/generative/layers/attention_utils.py,sha256=zBVwlBUTs-nStIKCZG0ks5ra7tsqc9ShfakFJKH5rds,7344
 ai_edge_torch/generative/layers/builder.py,sha256=LXGuSHIx6QZAzLFm7aJvlzoMPgQwbXLFchGEKYwOOUA,5090
 ai_edge_torch/generative/layers/feed_forward.py,sha256=hdICat-8gW7-vxDAevJQ8NQ-mynllPiqLdXQMF6JMnc,4189
-ai_edge_torch/generative/layers/kv_cache.py,sha256=9kkFpB9msgUDStFxEyQYYsavKPP4Dgqb_NFcd4hA4aU,8502
+ai_edge_torch/generative/layers/kv_cache.py,sha256=WNH_Ab29eXKXs8HAm3Wmdv_LBzO6PQW5d34Eo6Yzgd0,8492
 ai_edge_torch/generative/layers/lora.py,sha256=hsvWLLOnW7HQ0AysOZu30x_cetMquDd1tjfyLz8HCSU,17892
 ai_edge_torch/generative/layers/model_config.py,sha256=nLXvTkDAIHJQ0PTaWODF8oxJQoJ-K8D10cKR9229SAw,8355
 ai_edge_torch/generative/layers/normalization.py,sha256=MbwH-n80Fob5YvjBzdqDjBizMHLzSJGYRDdbD-rL5C0,6174
 ai_edge_torch/generative/layers/rotary_position_embedding.py,sha256=975zR202MdIrILJ7blceAcxrNqX1ZCN0ECKG1gz-bV8,2655
 ai_edge_torch/generative/layers/scaled_dot_product_attention.py,sha256=vp8dVx6tOe99neJhpbrtIt5fvN5NFw19JVH1v0yi5Mg,4154
-ai_edge_torch/generative/layers/sdpa_with_kv_update.py,sha256=TcwiI1IHhcYUrTx0kpSPAJMxFfjFcDwAHHULfZm67U4,3785
+ai_edge_torch/generative/layers/sdpa_with_kv_update.py,sha256=oo9h7pi0GcuylRgp2yUuvUJCrhj03aoWt_fP7EDP4LM,3775
 ai_edge_torch/generative/layers/experimental/__init__.py,sha256=nz-K0h8DfiATHzR6s1_bCw2akUmHWffU1bDRSkIzSqI,592
 ai_edge_torch/generative/layers/experimental/attention.py,sha256=XYbo1KlmiMEuwArye0Ul86jEsdxLr1RG-usRpidZiT8,8001
 ai_edge_torch/generative/layers/experimental/kv_cache.py,sha256=zgpFVftOfllvjh9-UEBSvUbm152SnQETn29rUMMMvAM,2978
-ai_edge_torch/generative/layers/experimental/scaled_dot_product_attention.py,sha256=8M6tC5kIUus-wbMEKDSMbCLnsobs6rgbujycsmhYa5g,2807
-ai_edge_torch/generative/layers/experimental/types.py,sha256=gZI9hIPB3XAo4oecKIIoVDfiyibLaSNFhecPFx4VDTM,2913
+ai_edge_torch/generative/layers/experimental/scaled_dot_product_attention.py,sha256=YFcIGOkaNb-vvQKjI-G9-bC2Z1W0O_qRyIZPlsLl72U,2797
 ai_edge_torch/generative/layers/unet/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/layers/unet/blocks_2d.py,sha256=ZteHZXK6HKyxYji49DQ46sA9aIy7U3Jnz0HZp6hfevY,28996
 ai_edge_torch/generative/layers/unet/builder.py,sha256=zAqWXdimmMrQRhmE_t9XkS68mh6PSrzwb-2NZZXrR5I,1901
@@ -195,6 +194,7 @@ ai_edge_torch/generative/utilities/moonshine_loader.py,sha256=_RpFabSqtGH5PHiP3_
 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=8sp9m_FMcXn7nqOrochtu2jIANkJKhnhIBUmH0ZTDR4,1549
+ai_edge_torch/generative/utilities/types.py,sha256=gZI9hIPB3XAo4oecKIIoVDfiyibLaSNFhecPFx4VDTM,2913
 ai_edge_torch/generative/utilities/verifier.py,sha256=gsBGv-WgeP73Bmao29CiTIQOC1YZ43IUJcGzytpcZyM,12095
 ai_edge_torch/hlfb/__init__.py,sha256=sH4um75na-O8tzxN6chFyp6Y4xnexsE7kUQpZySv6dE,735
 ai_edge_torch/hlfb/mark_pattern/__init__.py,sha256=JsVmYrM_JEuN_smMHXUsRlo3Liapp7UyktbPpPARwDk,5386
@@ -245,8 +245,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.5.0.dev20250410.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.5.0.dev20250410.dist-info/METADATA,sha256=8m6hxUmTT0arSNdEuo-mOyg1w9T3ekAvedvY2T6Opgw,2051
-ai_edge_torch_nightly-0.5.0.dev20250410.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.5.0.dev20250410.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.5.0.dev20250410.dist-info/RECORD,,
+ai_edge_torch_nightly-0.5.0.dev20250412.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.5.0.dev20250412.dist-info/METADATA,sha256=Q0Noz-4Joa8ZoVdxiBHapOXehAogQTUdAK4DtZlV4Tk,2051
+ai_edge_torch_nightly-0.5.0.dev20250412.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.5.0.dev20250412.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.5.0.dev20250412.dist-info/RECORD,,