remove the array with duplicate element

[yiminghwang]

Hi, there,

Is there any good way to remove the duplicated elements in DeviceArray?
we all know it can be implemented as

a_list = [[1,0,0,1],[1,0,1,1],[0,1,1,0],[1,0,0,1],[1,0,1,1]]
res_list = [list(t) for t in set(tuple(list(_)) for _ in a_list)]

Is there any effecient implementation using JAX? directly use fori_loop to filter?

[jakevdp]

You can do this with jnp.unique:

import jax.numpy as jnp
a = jnp.array([[1,0,0,1],[1,0,1,1],[0,1,1,0],[1,0,0,1],[1,0,1,1]])
print(jnp.unique(a, axis=0))
# [[0 1 1 0]
#  [1 0 0 1]
#  [1 0 1 1]]

[yiminghwang]

Oh! thanks, jake!