JAX: count unique elements in array

Question

I have a jitted function where one of the inputs is an array whose number of unique elements I would like to calculate. I cannot use jax.numpy.unique, it throws a ConcretizationTypeError due to the output of unique being of variable size. Any algorithmic workarounds_

Try using the size arg of jax.numpy.unique.

Chris Flesher
– Chris Flesher

2024-04-03 01:28:19 +00:00
Commented Apr 3, 2024 at 1:28 — Chris Flesher
– Chris Flesher, Commented Apr 3, 2024 at 1:28

jakevdp · Accepted Answer · 2023-09-12 13:50:33Z

2

One way to do this is to sort the array and count the number of adjacent entries which are not equal. For example:

import jax
import jax.numpy as jnp

x = jnp.array([3, 2, 1, 2, 3])

@jax.jit
def count_unique(x):
  x = jnp.sort(x)
  return 1 + (x[1:] != x[:-1]).sum()

print(count_unique(x))
# 3

answered Sep 12, 2023 at 13:50

jakevdp

88.6k15 gold badges140 silver badges177 bronze badges

Sign up to request clarification or add additional context in comments.

Collectives™ on Stack Overflow

JAX: count unique elements in array

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related