Repository: stanford-crfm/haliax (language: python)

Issue:
Worth adding top_k function?
It seems like adding a haliax impelementation for jax.lax.top_k would be a good idea, especially as it's being primarily used for language modeling by levanter. If this makes for a good first issue, I could maybe try my hand at it? The only thing that seems like it might be challenging is that you would have to resize the axis along which you are taking the top_k.


Discussion / hints:
sure, i'd take a PR for that!
Just to talk things through, here's where I'm at right now: 

1. My first thought was to try and wrap it using the existing wrapping functionality but this doesn't seem like it will work, because `jax.lax.top_k` returns two arrays, both an array of the top_k elements and their indices. Also, it doesn't fit neatly into a reduction or axiswise operation, since the axes remain the same, one is just resized. So it seems like it will need a custom wrapper, would you also agree?
2. One other thing is that `jax.lax.top_k` doesn't accept an axis to perform the operation on, it always chooses the last axis. Would the proper thing to do here be to allow the haliax version to accept an axis, and then transpose the underlying array such that the inputted axis becomes the last axis, and then perform the operation, and then reshape it back? 
3. Structurally, where does it belong in the codebase? Lots of similar-ish functions, e.g., hax.argmax are defined in the init file but none of these have custom wrappers. Would putting it in core.py (or somewhere else) make more sense?
(1) I think you'll have to use your own wrapper, yeah.

(2) yep, that would be great!

(3) maybe make a new file? I should move most of the things in init anyway



Fix patch:
diff --git a/src/haliax/__init__.py b/src/haliax/__init__.py
index 9ac54b6..97f4909 100644
--- a/src/haliax/__init__.py
+++ b/src/haliax/__init__.py
@@ -36,6 +36,7 @@
 from .hof import fold, map, scan, vmap
 from .ops import clip, isclose, pad_left, trace, tril, triu, where
 from .partitioning import auto_sharded, axis_mapping, fsdp, named_jit, shard_with_axis_mapping
+from .specialized_fns import top_k
 from .types import Scalar
 from .wrap import (
     ReductionFunction,
@@ -976,4 +977,5 @@ def true_divide(x1: NamedOrNumeric, x2: NamedOrNumeric, /) -> NamedOrNumeric:
     "selects_axis",
     "concat_axes",
     "concat_axis_specs",
+    "top_k",
 ]
diff --git a/src/haliax/specialized_fns.py b/src/haliax/specialized_fns.py
new file mode 100644
index 0000000..ec1008c
--- /dev/null
+++ b/src/haliax/specialized_fns.py
@@ -0,0 +1,25 @@
+from typing import Optional
+
+import jax
+import jax.numpy as jnp
+
+from .axis import Axis, AxisSelector
+from .core import NamedArray
+
+
+def top_k(arr: NamedArray, axis: AxisSelector, k: int, new_axis: Optional[AxisSelector] = None) -> NamedArray:
+    pos = arr._lookup_indices(axis)
+    if pos is None:
+        raise ValueError(f"Axis {axis} not found in {arr}")
+    new_array = jnp.moveaxis(arr.array, pos, -1)  # move axis to the last position
+    values, _ = jax.lax.top_k(new_array, k=k)
+    values = jnp.moveaxis(values, -1, pos)  # move axis back to its original position
+
+    if new_axis is None:
+        axis = arr.resolve_axis(axis)
+        new_axis = axis.resize(k)
+    elif isinstance(new_axis, str):
+        new_axis = Axis(new_axis, k)
+
+    updated_axes = arr.axes[:pos] + (new_axis,) + arr.axes[pos + 1 :]
+    return NamedArray(values, updated_axes)


Test patch:
diff --git a/tests/test_specialized_fns.py b/tests/test_specialized_fns.py
new file mode 100644
index 0000000..0e9c5a3
--- /dev/null
+++ b/tests/test_specialized_fns.py
@@ -0,0 +1,25 @@
+import jax
+import jax.numpy as jnp
+
+import haliax.specialized_fns as hfns
+from haliax import Axis, NamedArray
+
+
+def test_top_k():
+    H = Axis("H", 5)
+    W = Axis("W", 6)
+    D = Axis("D", 7)
+
+    rand = jax.random.uniform(jax.random.PRNGKey(0), (H.size, W.size, D.size))
+    n_rand = NamedArray(rand, (H, W, D))
+
+    assert hfns.top_k(n_rand, D, 2).array.shape == (H.size, W.size, 2)
+    assert jnp.all(
+        jnp.equal(jax.lax.top_k(rand, 2)[0], hfns.top_k(n_rand, D, 2).array)
+    )  # test that selecting last axis is same as default
+
+    for idx, i in enumerate(n_rand.axes):  # then test selecting all axes
+        t = jnp.transpose(rand, (*range(idx), *range(idx + 1, len(n_rand.axes)), idx))
+        t = jax.lax.top_k(t, 2)[0]
+        t = jnp.moveaxis(t, -1, idx)
+        assert jnp.all(jnp.equal(t, hfns.top_k(n_rand, i, 2).array))