WIP ENH: setdiff1d for Dask and jax.jit

[crusaderky]

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

[crusaderky]

@rgommers confirmed offline his preference for delaying indefinitely. His reasoning is that at some point JAX should support unknown shapes and the issue of niche functions that can't work until then should be marginal enough not to warrant urgent attention.

[lucascolley]

issue of niche functions that can't work until then should be marginal enough

Worth noting that this was hit for Dask in the cross-linked sklearn PR

[lucascolley]

@rgommers confirmed offline his preference for delaying indefinitely. His reasoning is that at some point JAX should support unknown shapes and the issue of niche functions that can't work until then should be marginal enough not to warrant urgent attention.

it would be nice to have an upstream issue at which to point, however 'at some point' it is