Follow-up to #115

• setdiff1d fails when executed inside jax.jit. This is because of a call to xp.unique* without an explicit size= parameter.
• By reading the code I'm fairly sure that the function fails if x1 and/or x2 have Dask backend and NaN in their shapes. Lacks a unit test.