Repository: jax-ml/jax
Issue category: Performance Issue

scipy.stats.rankdata causes constant folding warning for method='dense' but not method='ordinal'
### Description

[`scipy.stats.rankdata`](https://jax.readthedocs.io/en/latest/_autosummary/jax.scipy.stats.rankdata.html) causes a constant folding warning for `method='dense'` but not `method='ordinal'`:

```
$ py -c "import jax; jax.scipy.stats.rankdata(jax.numpy.zeros(10**7), 'ordinal')"
$ py -c "import jax; jax.scipy.stats.rankdata(jax.numpy.zeros(10**7), 'dense')"
2024-11-04 20:21:27.997499: E external/xla/xla/service/slow_operation_alarm.cc:65] Constant folding an instruction is taking > 1s:

  %reduce-window.6 = s32[625000,16]{0,1} reduce-window(s32[625000,16]{0,1} %constant.174, s32[] %constant.17), window={size=1x16 pad=0_0x15_0}, to_apply=%region_5.113

This isn't necessarily a bug; constant-folding is inherently a trade-off between compilation time and speed at runtime. XLA has some guards that attempt to keep constant folding from taking too long, but fundamentally you'll always be able to come up with an input program that takes a long time.

If you'd like to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.
2024-11-04 20:21:33.721446: E external/xla/xla/service/slow_operation_alarm.cc:133] The operation took 6.728977s
Constant folding an instruction is taking > 1s:

  %reduce-window.6 = s32[625000,16]{0,1} reduce-window(s32[625000,16]{0,1} %constant.174, s32[] %constant.17), window={size=1x16 pad=0_0x15_0}, to_apply=%region_5.113

This isn't necessarily a bug; constant-folding is inherently a trade-off between compilation time and speed at runtime. XLA has some guards that attempt to keep constant folding from taking too long, but fundamentally you'll always be able to come up with an input program that takes a long time.

If you'd like to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.
```

Looking at the code for `rankdata`, the culprit might be one of the 3 lines of code starting [here](https://github.com/jax-ml/jax/blob/ab47d4687f647de3aa145a9a782fb7b4aaf92af4/jax/_src/scipy/stats/_core.py#L206).

XLA dump [here](https://www.dropbox.com/scl/fo/rruuywlngh1r03hj9c2r1/AM-ym1pWfIUhkHA2hOiQNko?rlkey=2xxwdrmssgfyk7yz61xrt1t7d&st=h6yp3a8x&dl=0).

### System info (python version, jaxlib version, accelerator, etc.)

jax:    0.4.35
jaxlib: 0.4.34
numpy:  1.26.4
python: 3.12.7 (main, Oct  1 2024, 02:05:46) [Clang 15.0.0 (clang-1500.3.9.4)]
device info: cpu-1, 1 local devices"
process_count: 1
platform: uname_result(system='Darwin', node='Carloss-MacBook-Pro-2.local', release='23.6.0', version='Darwin Kernel Version 23.6.0: Mon Jul 29 21:14:46 PDT 2024; root:xnu-10063.141.2~1/RELEASE_ARM64_T6031', machine='arm64')


Functions to edit:
jax/_src/scipy/stats/_core.py:rankdata

Patch:
diff --git a/jax/_src/scipy/stats/_core.py b/jax/_src/scipy/stats/_core.py
index 08d1c0b6b538..f7b28d3ac301 100644
--- a/jax/_src/scipy/stats/_core.py
+++ b/jax/_src/scipy/stats/_core.py
@@ -198,13 +198,12 @@ def rankdata(
     return jnp.apply_along_axis(rankdata, axis, a, method)
 
   arr = jnp.ravel(a)
-  sorter = jnp.argsort(arr)
+  arr, sorter = jax.lax.sort_key_val(arr, jnp.arange(len(arr)))
   inv = invert_permutation(sorter)
 
   if method == "ordinal":
     return inv + 1
-  arr = arr[sorter]
-  obs = jnp.insert(arr[1:] != arr[:-1], 0, True)
+  obs = jnp.concatenate([jnp.array([True]), arr[1:] != arr[:-1]])
   dense = obs.cumsum()[inv]
   if method == "dense":
     return dense