ai-edge-torch-nightly 0.3.0.dev20250114__py3-none-any.whl → 0.3.0.dev20250116__py3-none-any.whl

ai_edge_torch/generative/layers/scaled_dot_product_attention.py

@@ -65,7 +65,7 @@ def scaled_dot_product_attention(
         scale=scale,
     )
   else:
-    q.mul_(scale)
+    q = q * scale
     scores = q @ k.transpose(-1, -2)
     scores = scores / softcap
     scores = torch.tanh(scores)
@@ -130,7 +130,7 @@ def scaled_dot_product_attention_with_hlfb(
         scale=scale,
     )
   else:
-    q.mul_(scale)
+    q = q * scale
     scores = q @ k.transpose(-1, -2)
     scores = scores / softcap
     scores = torch.tanh(scores)

ai_edge_torch/generative/utilities/converter.py

@@ -167,10 +167,7 @@ def _export_helper(
       prefill_seq_len = prefill_seq_lens[i]
       prefill_tokens = prefill_tokens_list[i]
       prefill_input_pos = prefill_input_pos_list[i]
-      if i == 0 and len(prefill_seq_lens) == 1:
-        prefill_signature_name = 'prefill'
-      else:
-        prefill_signature_name = f'prefill_{prefill_seq_len}'
+      prefill_signature_name = f'prefill_{prefill_seq_len}'
 
       sample_kwargs = {
           'tokens': prefill_tokens,

ai_edge_torch/version.py

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

{ai_edge_torch_nightly-0.3.0.dev20250114.dist-info → ai_edge_torch_nightly-0.3.0.dev20250116.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.3.0.dev20250114
+Version: 0.3.0.dev20250116
 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.3.0.dev20250114.dist-info → ai_edge_torch_nightly-0.3.0.dev20250116.dist-info}/RECORD

@@ -3,7 +3,7 @@ ai_edge_torch/_config.py,sha256=PKtOtBOup-cM0wBdQxby6HzuhLhIC3oq-TBG8FF4znE,2161
 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=aeyTcKMqD75DkncT9-cEtyswU4nSyQyc3hzDw71g_sU,706
+ai_edge_torch/version.py,sha256=G-QSz8n-nFkUu7vA6xPORcsrqKhus5kptuw7NkK8Au4,706
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=_PoH0E1gbbsWhLGwDRwUtW2G_IgNzNF7pKQbn9ct6-4,5778
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -127,7 +127,7 @@ ai_edge_torch/generative/layers/lora.py,sha256=hsvWLLOnW7HQ0AysOZu30x_cetMquDd1t
 ai_edge_torch/generative/layers/model_config.py,sha256=9yPEmWNw3-_2wXBmPmZ7RUKcPXHF2ZbJwksyQoXTA6M,7784
 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=gXxh3papKy4FBpGEX7VyZ7rZ1Js6aHK70Q6DKrVSckY,4154
+ai_edge_torch/generative/layers/scaled_dot_product_attention.py,sha256=vp8dVx6tOe99neJhpbrtIt5fvN5NFw19JVH1v0yi5Mg,4154
 ai_edge_torch/generative/layers/unet/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/generative/layers/unet/blocks_2d.py,sha256=9jKzOfiBQ66bp1ZnVIAoREIifVNFx4aTlQeYMAx2_pA,29062
 ai_edge_torch/generative/layers/unet/builder.py,sha256=zAqWXdimmMrQRhmE_t9XkS68mh6PSrzwb-2NZZXrR5I,1901
@@ -149,7 +149,7 @@ ai_edge_torch/generative/test/test_model_conversion_large.py,sha256=bBcey-aD4L_T
 ai_edge_torch/generative/test/test_quantize.py,sha256=bEJMhpQ9bIDUZVBXTW888728FcH-i3SyE4JSZZUgU0A,6071
 ai_edge_torch/generative/test/utils.py,sha256=tF6aCfAGJnc9dmzCnZCEOuKNVimfWOqscv9og0DDLHU,2656
 ai_edge_torch/generative/utilities/__init__.py,sha256=-_jxnnFnCgnTU4oTm4MnRsvL5lqhomBNdFBbqfmfHPo,720
-ai_edge_torch/generative/utilities/converter.py,sha256=MY8BK29yD-W4v45Xdl_ErbNilipsTlD-4-y9MyBxR5g,7620
+ai_edge_torch/generative/utilities/converter.py,sha256=yNIZ-O6RdXYl8yuWM_sTENRxozPnKGS-TZRhiiTaraE,7515
 ai_edge_torch/generative/utilities/dynamic_update_slice.py,sha256=e2mhx-Vp8sUK4EXoPtpZLSx3TViqLAKs67EhKcXBjAQ,2121
 ai_edge_torch/generative/utilities/loader.py,sha256=A3SOjPXp--AsvoP1hqj5QKWE4sgxoFc3H5EBUz_Eogc,13531
 ai_edge_torch/generative/utilities/model_builder.py,sha256=3CQLxJ02pFIo2DlS-RCn9cT6OvR4NiIuYRH597UXLiI,6530
@@ -206,8 +206,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.dev20250114.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.3.0.dev20250114.dist-info/METADATA,sha256=XzlB8Od_vqBIzPWNayCo-AAxty7ypWhNoQ0Ot7bJJ8w,1966
-ai_edge_torch_nightly-0.3.0.dev20250114.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.3.0.dev20250114.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.3.0.dev20250114.dist-info/RECORD,,
+ai_edge_torch_nightly-0.3.0.dev20250116.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.3.0.dev20250116.dist-info/METADATA,sha256=BBIWyQ45p6eKiiajHicFpac_J0nPnyAeviNewQm5_pc,1966
+ai_edge_torch_nightly-0.3.0.dev20250116.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.3.0.dev20250116.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.3.0.dev20250116.dist-info/RECORD,,