Closes #116
Needs more thorough unit tests + performance benchmarks.

This function's output is of unknown shape, so with the previous API it will never work in jax.jit.

There are a few options:

1. Stipulate that the function will never work inside jax.jit and you need to hack your way around it with ENH: lazy_apply #86 (comment).
   I'm not a fan of this because UX is very painful as it forces the user to think in graphs.
2. Add optional parameter fill_value. iff running inside the jax.jit, quietly return a longer array padded with it.
   I'm not happy about this because it causes jax.jit to quietly diverge from other backends and users will spend a lot of time debugging.
3. Add optional parameters size and fill_value. size becomes mandatory when running inside jax.jit. This is the same design as jax.numpy.unique_values.
   This also allows having a known-shape output in Dask. However, implementing it for Dask is fairly complicated.
4. Add optional parameters size and fill_value. size is mandatory when running inside jax.jit and disregarded otherwise. Again, this will cause bugs in the user code that only appear in jax.jit, but at least it demands an initial explicit user intervention. This is the simplest to implement; unsure on the UX. It also has the advantage of not sacrificing performance on other backends. If in the future jax.jit will support arrays of unknown size, it becomes easy to deprecate it as we said that the output size requested by the user may be disregarded anyway.

My current favourite is (4).
@rgommers you previously said, talking about functions with the same problem in scipy, that you prefer (1) to (3) because of not being able to retract the API in the future. What's your opinion on (4)?

CC @lucascolley