Simulation error with jax median operation about spu

[fionser]

Because secret sort is non-trivial. I'm also confused about it now, sort algorithm couldn't be parallelized. :(

try bitonic sort which is highly parallelable

[anakinxc]

Hi @BeStrongok,

jnp.median in this example requires secret sort, which is not supported by SPU right now.

But you can expect this feature in the upcoming releases.

Best

[BeStrongok]

Hi @BeStrongok,

jnp.median in this example requires secret sort, which is not supported by SPU right now.

But you can expect this feature in the upcoming releases.

Best

Thanks, looking for this feature!

[BeStrongok]

Because secret sort is non-trivial. I'm also confused about it now, sort algorithm couldn't be parallelized. :(

[anakinxc]

Because secret sort is non-trivial. I'm also confused about it now, sort algorithm couldn't be parallelized. :(

Yes, it is non-trivial and hard to optimize, but yes, we are working on it. :D
Stay tuned

[rivertalk]

Because secret sort is non-trivial. I'm also confused about it now, sort algorithm couldn't be parallelized. :(

try bitonic sort which is highly parallelable

Nice suggestion, trying...

[fionser]

Actually, current spu should be able to implement data structure like, top-k heap.
Then median can be done using top-(N/2) heap.

[anakinxc]

Actually, current spu should be able to implement data structure like, top-k heap. Then median can be done using top-(N/2) heap.

This really depends on how JAX implement median.
Right now they are using quantile(..., 0.5) to compute median, and quantile is using sort with a crazy comparator when input is floating point to handle nonfinite values and negative zero, which we don't really care.

During the compilation pipeline, when we have the semi-optimized IR, it is definitely possible to discover and rewrite this pattern with some efforts.

This kind of specific pattern rewrite is obviously not a unicorn, but still hacky...We are trying to see if this can be solved in a more elegant way.

[anakinxc]

Fixed with #107