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.
@@ -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
- key: val.func
47
- for key, val in torch_xla2.ops.ops_registry.all_aten_ops.items()
48
- if val.is_jax_function
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
@@ -13,4 +13,4 @@
13
13
  # limitations under the License.
14
14
  # ==============================================================================
15
15
 
16
- __version__ = "0.3.0.dev20241021"
16
+ __version__ = "0.3.0.dev20241023"
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: ai-edge-torch-nightly
3
- Version: 0.3.0.dev20241021
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=w4N5YyohM9rUsUUc1KQ3ViiuuI09nBTGGwwV0hL9ZcU,706
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=RN6BwMHuFj_rFgLCZ6Tu32XHbS2HGjPJeir2nROQ2rA,10517
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.dev20241021.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
185
- ai_edge_torch_nightly-0.3.0.dev20241021.dist-info/METADATA,sha256=QSx_Fkkgrk5_1B75iuQXcz1p07vsNPenHj59YBQmWuQ,1897
186
- ai_edge_torch_nightly-0.3.0.dev20241021.dist-info/WHEEL,sha256=eOLhNAGa2EW3wWl_TU484h7q1UNgy0JXjjoqKoxAAQc,92
187
- ai_edge_torch_nightly-0.3.0.dev20241021.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
188
- ai_edge_torch_nightly-0.3.0.dev20241021.dist-info/RECORD,,
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,,