@mmcky Can you please search for this line: In line with immutability, JAX does not support inplace operations: It looks like JAX is sorting in place now. Could you please evaluate and propose a fix?