How to jit trace jnp.histogramdd with different bin sizes on different dimensions?

[chjz1024]

I'm trying to jit trace a function that uses jnp.histogramdd in the computation. But wrapping a jax.jit around it raises an exception whenever the bins argument is a sequence of ints. Is it possible to perform jit tracing with different bin sizes on different dimensions? Assuming all subsequent calls will always use the same bins argument.

Example Code:

import jax
import jax.numpy as jnp
randarray = jax.random.normal(jax.random.key(0), (10000,3))
# jax.jit(jnp.histogramdd, static_argnames=['bins'])(randarray,bins=20) # works
# jax.jit(jnp.histogramdd, static_argnames=['bins'])(randarray,bins=[10,20,30]) # ValueError: Non-hashable static arguments are not supported. An error occurred while trying to hash an object of type <class 'list'>, [10, 20, 30]. The error was: TypeError: unhashable type: 'list'
# jax.jit(jnp.histogramdd)(randarray,bins=[10,20,30]) # ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: traced array with shape int32[] bins argument of histogram_bin_edges

[dfm]

There are a few options here. If you always want to use the same bins, something like:

from functools import partial

jax.jit(partial(jnp.histogramdd, bins=[10,20,30]))(randarray)

Another option would be to use a tuple instead of a list for your bins:

# jax.jit(jnp.histogramdd, static_argnames=['bins'])(randarray, bins=(10,20,30))

Hope this helps!

[chjz1024]

Thanks! The partial solution is very neat for arbitrary functions.