ai-edge-torch-nightly 0.3.0.dev20241101__py3-none-any.whl → 0.3.0.dev20241104__py3-none-any.whl

Sign up to get free protection for your applications and to get access to all the features.
@@ -63,10 +63,8 @@ _QUANTIZE = flags.DEFINE_bool(
63
63
 
64
64
  _DEVICE_TYPE = flags.DEFINE_string(
65
65
  'device_type',
66
- None,
66
+ 'cpu',
67
67
  help='The device type of the model. Currently supported: cpu, gpu.',
68
- default='cpu',
69
- required=True,
70
68
  )
71
69
 
72
70
 
@@ -19,9 +19,14 @@ from dataclasses import dataclass
19
19
  import gc
20
20
  import itertools
21
21
  import logging
22
+ import os
22
23
  import tempfile
23
24
  from typing import Any, Dict, Optional, Tuple, Union
24
25
 
26
+ if "PJRT_DEVICE" not in os.environ:
27
+ # https://github.com/google-ai-edge/ai-edge-torch/issues/326
28
+ os.environ["PJRT_DEVICE"] = "CPU"
29
+
25
30
  from ai_edge_torch import model
26
31
  from ai_edge_torch._convert import conversion_utils
27
32
  from ai_edge_torch._convert import signature as signature_module
ai_edge_torch/version.py CHANGED
@@ -13,4 +13,4 @@
13
13
  # limitations under the License.
14
14
  # ==============================================================================
15
15
 
16
- __version__ = "0.3.0.dev20241101"
16
+ __version__ = "0.3.0.dev20241104"
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: ai-edge-torch-nightly
3
- Version: 0.3.0.dev20241101
3
+ Version: 0.3.0.dev20241104
4
4
  Summary: Supporting PyTorch models with the Google AI Edge TFLite runtime.
5
5
  Home-page: https://github.com/google-ai-edge/ai-edge-torch
6
6
  Keywords: On-Device ML,AI,Google,TFLite,PyTorch,LLMs,GenAI
@@ -3,7 +3,7 @@ ai_edge_torch/config.py,sha256=FMWeCH2b7HYILBvaI1iZNnYCO4WAhDOwBZBmIE-xrF0,909
3
3
  ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
4
4
  ai_edge_torch/fx_pass_base.py,sha256=SrYveglaiA_DXPoRBqSXClWM1q7853I5ujRorq_MV0M,4251
5
5
  ai_edge_torch/model.py,sha256=N-pNpTxzhaFGhWhnSGd70lBzb9VlEhTOq5mddU7bvvI,5542
6
- ai_edge_torch/version.py,sha256=qcmpL1uPmUip6z4mTeUEry2vsiNcJ_XR-psTiLB1uo4,706
6
+ ai_edge_torch/version.py,sha256=dJo-nYwMDCrbqPdvTgYKKdINoavDEmUZMzMQWTPxUKA,706
7
7
  ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
8
8
  ai_edge_torch/_convert/conversion.py,sha256=HwzfRx_DX5TLtPqwEH1_NOm38_INvHzHl4_mX67KOdQ,5448
9
9
  ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -78,7 +78,7 @@ ai_edge_torch/generative/examples/smollm/verify.py,sha256=HXYcCjDJMylVL3Pc9HU-UX
78
78
  ai_edge_torch/generative/examples/stable_diffusion/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
79
79
  ai_edge_torch/generative/examples/stable_diffusion/attention.py,sha256=kDWG6MlIGa89zC5KSRcJlw2c4ITuw8KcchtfmF55f4g,3545
80
80
  ai_edge_torch/generative/examples/stable_diffusion/clip.py,sha256=5M4auM33SgCTODt0VT8TO-EVILruqGDRiNILBPeB83Y,6072
81
- ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py,sha256=Fw0ZsJSG8fM-07mEi6QLCn6LpRveGGL8vt7OIn0Av5c,5276
81
+ ai_edge_torch/generative/examples/stable_diffusion/convert_to_tflite.py,sha256=sB_7-PVri8PxKnFG7c8GsTGyrxGEda-oZwGyyScTL3o,5239
82
82
  ai_edge_torch/generative/examples/stable_diffusion/decoder.py,sha256=sQKQ-k6H9kG2brgwLsktjCMeN2h0POyfMP6iNsPNKWc,16271
83
83
  ai_edge_torch/generative/examples/stable_diffusion/diffusion.py,sha256=6W58LxmHHkz2ctgpknQkyoDANZAnE9Byp_svfqLpQf0,34793
84
84
  ai_edge_torch/generative/examples/stable_diffusion/encoder.py,sha256=CAPsW84A8f00nS6fLFeh_XUjCPsDCA5UxHOUsMrLfSU,3450
@@ -152,7 +152,7 @@ ai_edge_torch/lowertools/_shim.py,sha256=ilL7x1ebUBj1clg7bagrX4y_nVSHiGrvDrOVfuT
152
152
  ai_edge_torch/lowertools/common_utils.py,sha256=Z7p-ivOHtddktpnHrlDm_dSoTxJOdEjFXIGQbzjgwQo,4504
153
153
  ai_edge_torch/lowertools/odml_torch_utils.py,sha256=K5dZ_fFDL3GWKo0IoY4OC_GX5MY-guY-MqteolyV9hg,8098
154
154
  ai_edge_torch/lowertools/test_utils.py,sha256=bPgc2iXX16KYtMNvmsRdKfrCY6UJmcfitfCOvHoD7Oc,1930
155
- ai_edge_torch/lowertools/torch_xla_utils.py,sha256=S7RWzauts-15xP6VYuM3aAd9cyAGHstYD2A4dlv3d30,9059
155
+ ai_edge_torch/lowertools/torch_xla_utils.py,sha256=XGZE0vZG9WSQT-6dFmPlU8W89z8rfXPRGjuZeuhXCIw,9205
156
156
  ai_edge_torch/lowertools/translate_recipe.py,sha256=ymkBpFqAUiupRWqrPOWiVphKcXR1K5vHK0RjgBFtxlE,5652
157
157
  ai_edge_torch/odml_torch/__init__.py,sha256=S8jOzE9nLof-6es3XDiGJRN-9H_XTxsVm9dE7lD3RWo,812
158
158
  ai_edge_torch/odml_torch/_torch_future.py,sha256=jSYHf1CMTJzMizPMbu2b39hAt0ZTR6gQLq67GMe9KTo,2336
@@ -186,8 +186,8 @@ ai_edge_torch/quantize/quant_config.py,sha256=U0KisSW-uZkoMJcy-ZP9W57p3tsa594fr9
186
186
  ai_edge_torch/testing/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
187
187
  ai_edge_torch/testing/model_coverage/__init__.py,sha256=5P8J6Zk5YYtDvTBucFvB9NGSRI7Gw_24WnrbhXgycEE,765
188
188
  ai_edge_torch/testing/model_coverage/model_coverage.py,sha256=UPB448aMDUyC0HNYVqio2rcJPnDN0tBQMP08J6vPYew,4718
189
- ai_edge_torch_nightly-0.3.0.dev20241101.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
190
- ai_edge_torch_nightly-0.3.0.dev20241101.dist-info/METADATA,sha256=kxEaOa7LhiE2RVn3qSRQ0fRDKJgkdF908tN3QWlXnMY,1897
191
- ai_edge_torch_nightly-0.3.0.dev20241101.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
192
- ai_edge_torch_nightly-0.3.0.dev20241101.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
193
- ai_edge_torch_nightly-0.3.0.dev20241101.dist-info/RECORD,,
189
+ ai_edge_torch_nightly-0.3.0.dev20241104.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
190
+ ai_edge_torch_nightly-0.3.0.dev20241104.dist-info/METADATA,sha256=5KCg_KPNSulWAdcVRe-TY_L1WwdOzISVDWJ8eKc_saQ,1897
191
+ ai_edge_torch_nightly-0.3.0.dev20241104.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
192
+ ai_edge_torch_nightly-0.3.0.dev20241104.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
193
+ ai_edge_torch_nightly-0.3.0.dev20241104.dist-info/RECORD,,