ai-edge-torch-nightly 0.7.0.dev20250915__py3-none-any.whl → 0.7.0.dev20250917__py3-none-any.whl

ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py

@@ -18,6 +18,7 @@ import logging
 from ai_edge_torch.odml_torch import jax_bridge
 from ai_edge_torch.odml_torch.lowerings import context
 from ai_edge_torch.odml_torch.lowerings import registry
+import jax
 import jax.numpy as jnp
 from jax._src.lib.mlir import ir
 import torch
@@ -218,7 +219,6 @@ lower_by_torch_xla2(torch.ops.aten.tensor_split.sections)
 lower_by_torch_xla2(torch.ops.aten.to.device)
 lower_by_torch_xla2(torch.ops.aten.to.device)
 lower_by_torch_xla2(torch.ops.aten.to.dtype)
-lower_by_torch_xla2(torch.ops.aten.topk)
 lower_by_torch_xla2(torch.ops.aten.transpose)
 lower_by_torch_xla2(torch.ops.aten.transpose_copy)
 lower_by_torch_xla2(torch.ops.aten.triu)
@@ -508,3 +508,36 @@ def _aten_einsum_default(
     return jnp.einsum(equation, *operands, optimize="optimal")
 
   return jax_lowering(lctx, tuple(tensors))
+
+
+@registry.lower(torch.ops.aten.topk)
+def _aten_topk(
+    lctx: LoweringContext, self, k, dim=-1, largest=True, sorted=True
+):
+  _log_usage(torch.ops.aten.topk)
+
+  if not sorted:
+    logging.warning(
+        "aten.topk lowering ignores `sorted=False` and always returns sorted"
+        " results."
+    )
+
+  @jax_bridge.wrap
+  def jax_lowering(self, k):
+    if not largest:
+      self = -self
+    # jax.lax.top_k always sorts and operates on the last dimension.
+    move_dim_to_last = dim != -1 and dim != self.ndim - 1
+    if move_dim_to_last:
+      input_tensor = jnp.moveaxis(self, dim, -1)
+    else:
+      input_tensor = self
+    values, indices = jax.lax.top_k(input_tensor, k)
+    if move_dim_to_last:
+      values = jnp.moveaxis(values, -1, dim)
+      indices = jnp.moveaxis(indices, -1, dim)
+    if not largest:
+      values = -values
+    return values, indices
+
+  return jax_lowering(lctx, self, k)

ai_edge_torch/version.py

@@ -15,4 +15,4 @@
 
 # The next version of ai-edge-torch.
 # The minor version code should be bumped after every release.
-__version__ = "0.7.0.dev20250915"
+__version__ = "0.7.0.dev20250917"

{ai_edge_torch_nightly-0.7.0.dev20250915.dist-info → ai_edge_torch_nightly-0.7.0.dev20250917.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ai-edge-torch-nightly
-Version: 0.7.0.dev20250915
+Version: 0.7.0.dev20250917
 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.7.0.dev20250915.dist-info → ai_edge_torch_nightly-0.7.0.dev20250917.dist-info}/RECORD

@@ -2,7 +2,7 @@ ai_edge_torch/__init__.py,sha256=lemyLCNoGYRnJsmDuGZu7qOqLbLqG6CGDFtu3ue1syU,129
 ai_edge_torch/_config.py,sha256=AiqhbcheF7j_ozIGDLC89k1we95aVgFDa-tR6h7UI0s,2529
 ai_edge_torch/conftest.py,sha256=r0GTrhMRhlmOGrrkvumHN8hkmyug6WvF60vWq8wRIBI,758
 ai_edge_torch/model.py,sha256=A7loFu8jE9CsXsfMmHYZ-KDFJiaD8Kkqwm_9d3IVzk0,5638
-ai_edge_torch/version.py,sha256=sj61xhc7PcSll3exrFNmhToMA_lBgYwQ1IMBgahQ9rQ,806
+ai_edge_torch/version.py,sha256=RJAGmMeczkThv7SjCLM7f3fY57qXGMf2nQcJkjQqyB4,806
 ai_edge_torch/_convert/__init__.py,sha256=hHLluseD2R0Hh4W6XZRIXY_dRQeYudjsrKGf6LZz65g,671
 ai_edge_torch/_convert/conversion.py,sha256=iQk3R-pLq4c1nfLqPB4xTRj78gghxPGzJCJtILLdg5o,6123
 ai_edge_torch/_convert/conversion_utils.py,sha256=Sr8qXVcTwc-ZnZmK7yxVrIOOp1S_vNrwzC0zUvLTI2o,2160
@@ -253,7 +253,7 @@ ai_edge_torch/odml_torch/lowerings/_basic.py,sha256=sC4N5-7RS9yKecs97kM9J56enGvs
 ai_edge_torch/odml_torch/lowerings/_batch_norm.py,sha256=PaLI0BB6pdBW1VyfW8VTOT_Be-ZcqYdNOsyfzKfq8Cg,2064
 ai_edge_torch/odml_torch/lowerings/_convolution.py,sha256=Q0aDzyUcZMoSzSbOU-r3LJMgPe6fble0QwdYVIOHHHk,6887
 ai_edge_torch/odml_torch/lowerings/_decomp_registry.py,sha256=ybOdoFE5HIJTkyiYcc73zpyUyUpioVnAca6k0wyJPs4,2572
-ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=Hth8F1SwPVEympR1_T6sEsCmmRXtoO7zatPR1uRLFbY,18411
+ai_edge_torch/odml_torch/lowerings/_jax_lowerings.py,sha256=6KZc0-SM0LupOnDAgnUjj47aHk7YJ0cC1ulAeWvV5_w,19289
 ai_edge_torch/odml_torch/lowerings/_layer_norm.py,sha256=khJIvDVk2s332Nd2Be-5dM6-wp5DGff61HCV5lskHmQ,3011
 ai_edge_torch/odml_torch/lowerings/_quantized_decomposed.py,sha256=XDZ0zLej_XaQDJnaAAxhNFAd7NfQm5SOVEp_nno_krA,6178
 ai_edge_torch/odml_torch/lowerings/_rand.py,sha256=g6SuqDkuC6hD35lyP1-5H7ASDIzPSmKukeNT5naZSv8,4133
@@ -269,8 +269,8 @@ ai_edge_torch/testing/__init__.py,sha256=_yGgvnBZWb7T3IN3mc4x1sS4vM96HZwM8pwIcPG
 ai_edge_torch/testing/export.py,sha256=k5mGDGzwc23Z4zaIVDs8CNh-oOt64gsf9MS9NjhbPy4,3293
 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.7.0.dev20250915.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
-ai_edge_torch_nightly-0.7.0.dev20250915.dist-info/METADATA,sha256=jSrydAnNyQUAsVL_K2iue8o-seM1IevQdaK5NwDNlAc,2074
-ai_edge_torch_nightly-0.7.0.dev20250915.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-ai_edge_torch_nightly-0.7.0.dev20250915.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
-ai_edge_torch_nightly-0.7.0.dev20250915.dist-info/RECORD,,
+ai_edge_torch_nightly-0.7.0.dev20250917.dist-info/LICENSE,sha256=z8d0m5b2O9McPEK1xHG_dWgUBT6EfBDz6wA0F7xSPTA,11358
+ai_edge_torch_nightly-0.7.0.dev20250917.dist-info/METADATA,sha256=-mF2daWYgvauqMwRg8-wNaVQZ2hGtTRBrjPql34gNIU,2074
+ai_edge_torch_nightly-0.7.0.dev20250917.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+ai_edge_torch_nightly-0.7.0.dev20250917.dist-info/top_level.txt,sha256=5KXRaF2hwkApYxf7Y8y_tVb9aulGTlbOoNdbx1aKRkE,14
+ai_edge_torch_nightly-0.7.0.dev20250917.dist-info/RECORD,,