How to implement an JIT compatible groupby?

[n2cholas]

I'd like to implement a generic group-by then aggregate in JAX which has similar semantics to the function below:

def groupby_agg(keys, vals, agg_fn):
    return jnp.stack([agg_fn(vals[keys==i]) for i in jnp.unique(keys)])

keys = jnp.array([0, 1, 0, 2, 1])
vals = jnp.array([1.1, 1.2, 2.1, 2.2, 3.1])
groupby_agg(keys, vals, agg_fn=jnp.mean)
# DeviceArray([1.5999999, 2.15     , 2.2      ], dtype=float32)

JAX doesn't support dynamic sizes or advanced boolean indexing under JIT, so I thought something like this might work:

def groupby_agg2(keys, vals, unique_keys, agg_fn):
    def helper(i):
        arr = jnp.where(keys == i, vals, jnp.inf).sort()
        return agg_fn(jax.lax.dynamic_slice(arr, [0], [arr.searchsorted(jnp.inf)]))
    return jax.vmap(helper)(unique_keys)

groupby_agg2(keys, vals, jnp.unique(keys), agg_fn=jnp.mean)

But jax.lax.dynamic_slice requires slice_sizes to be static. I could not find a way to get this to work for a generic agg_fn. For some functions, such as a mean, it can be done simply:

def groupby_mean(keys, vals, unique_keys):
    def helper(i):
        mask = keys == i
        return jnp.where(mask, vals, 0).sum() / mask.sum()
    return jax.vmap(helper)(unique_keys)

groupby_mean(keys, vals, jnp.unique(keys))
# DeviceArray([1.5999999, 2.15     , 2.2      ], dtype=float32)

However, it's harder for more complicated aggregation functions since finding a "filler" value for jnp.where can be non-trivial.

Are there any dynamic masking functions available in JAX? If not, is there any other way to accomplish this? Thanks in advance!

[Joshuaalbert]

For a summing aggregator you can use bin_count and the weights argument. You do need to pass what the maximum number of unique keys is, and to encode the keys as ints. Does that do it?

[n2cholas]

Thanks, that's helpful in some cases, but not in the general case, which I'm interested in. For example, that won't work for "max".

@jakevdp has a nice blog post with ways to implement groupby in raw numpy: https://jakevdp.github.io/blog/2017/03/22/group-by-from-scratch/, but these don't help so much for the general JIT case.

[jakevdp]

Here's a quick groupby summation implementation in JAX:

import jax.numpy as jnp
import jax

def jax_groupby(keys, vals, unique_keys=None):
  if unique_keys is None:
    unique_keys, group = jnp.unique(keys, return_inverse=True)
  else:
    group = jnp.searchsorted(unique_keys, keys, side='left')
  counts = jnp.bincount(group, weights=vals, length=len(unique_keys) + 1)[:-1]
  return unique_keys, counts

If you pass a sorted list of unique_keys to the function, it will even be compatible with JIT, I believe.

A general groupby that is compatible with JIT will have to depend on something like masking support for the value-dependent output shapes.

[n2cholas]

Thanks! Since my request isn't possible with the current library, I'll mark it as answered.