ai-edge-torch-nightly 0.3.0.dev20241021__py3-none-any.whl → 0.3.0.dev20241023__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py +17 -6
- ai_edge_torch/version.py +1 -1
- {ai_edge_torch_nightly-0.3.0.dev20241021.dist-info → ai_edge_torch_nightly-0.3.0.dev20241023.dist-info}/METADATA +1 -1
- {ai_edge_torch_nightly-0.3.0.dev20241021.dist-info → ai_edge_torch_nightly-0.3.0.dev20241023.dist-info}/RECORD +7 -7
- {ai_edge_torch_nightly-0.3.0.dev20241021.dist-info → ai_edge_torch_nightly-0.3.0.dev20241023.dist-info}/LICENSE +0 -0
- {ai_edge_torch_nightly-0.3.0.dev20241021.dist-info → ai_edge_torch_nightly-0.3.0.dev20241023.dist-info}/WHEEL +0 -0
- {ai_edge_torch_nightly-0.3.0.dev20241021.dist-info → ai_edge_torch_nightly-0.3.0.dev20241023.dist-info}/top_level.txt +0 -0
@@ -42,11 +42,18 @@ def lower_by_jax(op, ir_input_names=None):
|
|
42
42
|
return inner
|
43
43
|
|
44
44
|
|
45
|
-
_TORCH_XLA2_IMPLS = {
|
46
|
-
|
47
|
-
|
48
|
-
|
49
|
-
|
45
|
+
_TORCH_XLA2_IMPLS = {}
|
46
|
+
|
47
|
+
for op, torch_xla2_op in torch_xla2.ops.ops_registry.all_aten_ops.items():
|
48
|
+
if not torch_xla2_op.is_jax_function:
|
49
|
+
continue
|
50
|
+
if isinstance(op, torch._ops.OpOverloadPacket):
|
51
|
+
ops = [getattr(op, overload) for overload in op.overloads()] + [op]
|
52
|
+
else:
|
53
|
+
ops = [op]
|
54
|
+
|
55
|
+
for op in ops:
|
56
|
+
_TORCH_XLA2_IMPLS[op] = torch_xla2_op.func
|
50
57
|
|
51
58
|
|
52
59
|
def lower_by_torch_xla2(op):
|
@@ -229,7 +236,6 @@ lower_by_torch_xla2(torch.ops.aten.transpose_copy)
|
|
229
236
|
lower_by_torch_xla2(torch.ops.aten.triu)
|
230
237
|
lower_by_torch_xla2(torch.ops.aten.true_divide)
|
231
238
|
lower_by_torch_xla2(torch.ops.aten.trunc)
|
232
|
-
lower_by_torch_xla2(torch.ops.aten.unbind)
|
233
239
|
lower_by_torch_xla2(torch.ops.aten.unbind_copy)
|
234
240
|
lower_by_torch_xla2(torch.ops.aten.unsqueeze)
|
235
241
|
lower_by_torch_xla2(torch.ops.aten.unsqueeze.default)
|
@@ -247,6 +253,11 @@ lower_by_torch_xla2(torch.ops.prims.broadcast_in_dim)
|
|
247
253
|
lower_by_torch_xla2(torch.ops.prims.var)
|
248
254
|
|
249
255
|
|
256
|
+
@lower_by_jax(torch.ops.aten.unbind)
|
257
|
+
def _aten_copy(self, *args, **kwargs):
|
258
|
+
return _TORCH_XLA2_IMPLS[torch.ops.aten.unbind_copy](self, *args, **kwargs)
|
259
|
+
|
260
|
+
|
250
261
|
@lower_by_jax(torch.ops.aten.copy, ir_input_names=["src"])
|
251
262
|
def _aten_copy(self, src, **kwargs):
|
252
263
|
return _TORCH_XLA2_IMPLS[torch.ops.aten.copy](self, src)
|
ai_edge_torch/version.py
CHANGED
@@ -1,6 +1,6 @@
|
|
1
1
|
Metadata-Version: 2.1
|
2
2
|
Name: ai-edge-torch-nightly
|
3
|
-
Version: 0.3.0.
|
3
|
+
Version: 0.3.0.dev20241023
|
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=
|
6
|
+
ai_edge_torch/version.py,sha256=GjmtIqPIRJ9PpLDTs_m_Jbpvogx7PJ-OK8Bhk-TfbOw,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
|
@@ -168,7 +168,7 @@ ai_edge_torch/odml_torch/lowerings/__init__.py,sha256=dE_qzh-OnCNjWzqs1-PHs5PNlR
|
|
168
168
|
ai_edge_torch/odml_torch/lowerings/_basic.py,sha256=VvB050UCjB17h6-UNtsaqzVF13MGI01fPFkdmmghTj4,8790
|
169
169
|
ai_edge_torch/odml_torch/lowerings/_batch_norm.py,sha256=PaLI0BB6pdBW1VyfW8VTOT_Be-ZcqYdNOsyfzKfq8Cg,2064
|
170
170
|
ai_edge_torch/odml_torch/lowerings/_convolution.py,sha256=v1VdKmL8YLJv3PR9VgyNghO83A25PpTzY2ZUAJqlq3Q,6847
|
171
|
-
ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=
|
171
|
+
ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=UI_2bbtbkZLBVVcYBXvA6JT1Fg2W_E_OSKYHT-D9xXI,10854
|
172
172
|
ai_edge_torch/odml_torch/lowerings/_layer_norm.py,sha256=1ePJs7oIdUkVdMddFsXMc53qTkEKqGz0ZhQQoNzBa10,2862
|
173
173
|
ai_edge_torch/odml_torch/lowerings/context.py,sha256=jslcCv7r_HtImSRTxJwHAUV_QCu9Jub51lovmoBkmFA,1295
|
174
174
|
ai_edge_torch/odml_torch/lowerings/registry.py,sha256=gqx3n1Mx8pnGQz3nkIF1T_8bkRabXLJBvUoJJn5kOUY,2911
|
@@ -181,8 +181,8 @@ ai_edge_torch/quantize/quant_config.py,sha256=U0KisSW-uZkoMJcy-ZP9W57p3tsa594fr9
|
|
181
181
|
ai_edge_torch/testing/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
|
182
182
|
ai_edge_torch/testing/model_coverage/__init__.py,sha256=5P8J6Zk5YYtDvTBucFvB9NGSRI7Gw_24WnrbhXgycEE,765
|
183
183
|
ai_edge_torch/testing/model_coverage/model_coverage.py,sha256=UPB448aMDUyC0HNYVqio2rcJPnDN0tBQMP08J6vPYew,4718
|
184
|
-
ai_edge_torch_nightly-0.3.0.
|
185
|
-
ai_edge_torch_nightly-0.3.0.
|
186
|
-
ai_edge_torch_nightly-0.3.0.
|
187
|
-
ai_edge_torch_nightly-0.3.0.
|
188
|
-
ai_edge_torch_nightly-0.3.0.
|
184
|
+
ai_edge_torch_nightly-0.3.0.dev20241023.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
|
185
|
+
ai_edge_torch_nightly-0.3.0.dev20241023.dist-info/METADATA,sha256=OfznHDTn9U_zuVoiqlyRXD4CngeLHzKQXKCaXxL8IXs,1897
|
186
|
+
ai_edge_torch_nightly-0.3.0.dev20241023.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
|
187
|
+
ai_edge_torch_nightly-0.3.0.dev20241023.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
|
188
|
+
ai_edge_torch_nightly-0.3.0.dev20241023.dist-info/RECORD,,
|
File without changes
|
File without changes
|