[torchax] Fixes test for kthvalue #7458 (#9223)

Author: qianminj123

torchax/test/test_ops.py

@@ -19,7 +19,6 @@
     "histogram",  # hard op: AssertionError: Tensor-likes are not close!
     "histogramdd",  # TypeError: histogram requires ndarray or scalar arguments, got <class 'list'> at position 1.
     "index_reduce",
-    "kthvalue",
     "linalg.ldl_solve",
     "max_pool2d_with_indices_backward",
     "nn.functional.adaptive_max_pool1d",
@@ -175,7 +174,7 @@ def run_export_and_compare(testcase,
 # Sort related ops should ignore index;
 # For example: sort( [1, 0, 0]) -> [0, 0, 1]
 # the correct index can be [1, 2, 0] or [2, 1, 0]
-should_ignore_indexes = {"topk", "mode"}
+should_ignore_indexes = {"topk", "mode", "kthvalue"}
 
 
 class TestOpInfo(TestCase):

torchax/torchax/ops/jaten.py

@@ -5496,3 +5496,20 @@ def linear(input, weight, bias=None):
   if bias is not None:
     res += bias
   return res
+
+
+@op(torch.ops.aten.kthvalue)
+def kthvalue(input, k, dim=None, keepdim=False, *, out=None):
+  if input.ndim == 0:
+    return input, jnp.array(0)
+  dimension = -1
+  if dim is not None:
+    dimension = dim
+  while dimension < 0:
+    dimension = dimension + input.ndim
+  values = jax.lax.index_in_dim(
+      jnp.partition(input, k - 1, dimension), k - 1, dimension, keepdim)
+  indices = jax.lax.index_in_dim(
+      jnp.argpartition(input, k - 1, dimension).astype('int64'), k - 1,
+      dimension, keepdim)
+  return values, indices