The following code does not work:
def get_unique(arr):
return jnp.unique(arr)
get_unique = jit(get_unique)
get_unique(jnp.ones((10,)))
The error message compains about the use of jnp.unique:
FilteredStackTrace: jax._src.errors.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(float32[10])>with<DynamicJaxprTrace(level=0/1)>
The error arose in jnp.unique()
The documentation on sharp bits explains that jit doesn't work if the shape of internal arrays depends on argument values. This is exactly the case here.
According to the docs, a potential workaround is to specify static parameters. But this doesn't apply to my case. The parameters will change for almost every function call. I have split up my code into a preprocessing step, which performs calculations such as this jnp.unique, and a computation step which can be jitted.
But still I'd like to ask, is there some workaround that I'm not aware of?
