scipy.stats.mode

[gianlucadetommaso]

It would be great if you could add the JAX equivalent of scipy.stats.mode, which is currently unavailable.

A use case may be a ML classification task with ensembles, where multiple models give different predictions, and we are interested in finding the most common one.

As an example, consider predictions to be a two-dimensional DeviceArray of predictions, with
shape = (number of models, number of data points).

It would be nice if one could compute

mode_prediction = jax.scipy.stats.mode(predictions, axis=0),

returning the most common prediction for every data point.

[jakevdp]

Here's a quick jit-compatible implementation if you want to make use of it (not thoroughly tested):

import jax.numpy as jnp
from jax import vmap

def mode(x, axis=None):
  def _mode(x):
    vals, counts = jnp.unique(x, return_counts=True, size=x.size)
    return vals[jnp.argmax(counts)]
  if axis is None:
    return _mode(x)
  else:
    x = jnp.moveaxis(x, axis, 0)
    return vmap(_mode, in_axes=(1,))(x.reshape(x.shape[0], -1)).reshape(x.shape[1:])

x = jnp.array([1, 1, 2, 2, 2, 3, 4])
print(mode(x))
# 2

y = jnp.array([[1, 1, 2, 3],
               [4, 5, 6, 5]])
print(mode(y, axis=1))
# [1, 5]

[gianlucadetommaso]

Great, thanks!

[jakevdp]

It's definitely in scope to add this to jax.scipy.stats. If anyone wants to tackle it, you could take the example implementation from above, modify it to make sure it returns the same outputs/shapes as scipy.stats (I think scipy.stats.mode returns the counts as well), and add an appropriate test in tests/scipy_stats_test.py

[jayendra13]

Let me take this up, can anyone suggest a good location/file for this function to put ?
For Scipy it's defined here and here.

I am starting this for now by putting it inside the init.py

[Yash621]

@gianlucadetommaso can i work on this ?

[gianlucadetommaso]

@Yash621 please go ahead, I haven't had much time to take this up myself.

[Yash621]

@gianlucadetommaso can you suggest the location where i should place this function?

[jakevdp]

You could put the implementation in a new submodule, maybe something like jax/_src/scipy/stats/_core.py

[elusenji]

I'd love to jump in and collaborate with @Yash621 on this, if that's okay.

[parthpatel9414]

Would love to collaborate with @Yash621 and @elusenji, if this issue is still open?

[jakevdp]

I think the PR in #9141 is quite close - it just needs a bit more work before it can be merged.