How to compute top_k and return k-hot encoded vector

[BugQualia]

tl;dr: I want to write this top_k_hot function which is 1. jittable, 2. minimal performance loss:

which takes two input:

1. arr: array of size n
2. k: integer k

and outputs:
array of size n which is marked 1 in the top-k index of arr
marked 0 if its not in the top-k index of arr

simple example:

arr = jnp.array([1, 6, 3, 4, 5], jnp.float32)
top_k_hot(arr, 2)
# output: jnp.array([0, 1, 0, 0, 1])

---

A simple implementation would be to use jax.lax.top_k.
However, this is not jit-able because k is variable,
and jax.lax.top_k outputs the array of size k.

def top_k_hot_not_jittable(arr, k):
    val, ind = jax.lax.top_k(arr, k)  # val.size == ind.size == k
    ret = jnp.zeros(arr.shape[0], jnp.int32)
    ret = jax.ops.index_update(ret, ind, 1)
    return ret

aa = jnp.array([10, 3, 5, 7, 9])
top_k_hot_not_jittable(aa, 3)

# DeviceArray([1, 0, 0, 1, 1], dtype=int32)

There are some alternatives I came up with, but it still doesn't solves the problem that
jax.lax.top_k outputs variable-sized array.
Arg-sorting the array and using the sliced index also doesn't work because jax.ops.index_update should not take variable
start:stop:step index to be jit-able.
Sorting the array and taking the right value works, but it acts differently and I'm quite worried about the performance.

@jax.jit
def top_k_hot(arr, k):
    arr_sorted = jnp.sort(arr)
    top_k_val = jnp.take(arr_sorted, arr.size-k-1, axis=0)
    return arr > top_k_val

aa = jnp.array([10, 3, 5, 7, 9])
top_k_hot(aa, 3)

# DeviceArray([ True, False, False,  True,  True], dtype=bool)

Partial sorting could improve a performance a little.
But this function will execute so frequently, I want optimal solution if possible.

So it seems like one can implement jit-able top_k_hot function,
but I can only do it with sad performance loss for now.

Any help will be appreciated!!

---

Also, thanks for making this awesome library.
It feels like I've finally found a right tool to use.

[cgarciae]

Why not set k as a static argument to the first function?

@partial(jax.jit, static_argnums=(1,))
def top_k_hot(arr, k):
    val, ind = jax.lax.top_k(arr, k)  # val.size == ind.size == k
    ret = jnp.zeros(arr.shape[0], jnp.int32)
    ret = jax.ops.index_update(ret, ind, 1)
    return ret

Been trying to do something similar but with 2D arrays and index_update was giving odd results so I resorted to this:

    ...
    val, idx = jax.lax.top_k(prob_tensor, topk)
    tops = jax.nn.one_hot(idx, prob_tensor.shape[-1], dtype=jnp.int32)
    return tops.sum(axis=-2)