"ValueError: Non-hashable static arguments" when vectorizing for loop

[guglielmogattiglio]

Hi, I have a function which does something and it's behavior depends on a (static) number. For example we can
consider a sum between shifted arrays (original array and it's s-lag)

import jax
import jax.numpy as jnp

def shifted_sum(x, y, s):
    return jnp.sum(x[s:] + y[:-s])

Say that I am now interested in the following (sum of shifted sums)

tot = 0
S = 4
x = y = jnp.arange(5)
for s in range(1, S):
    tot += shifted_sum(x, y, s)
print(tot)

Since we are eventually summing the output of shifted_sum(), the whole loop should be easily parallelizable in jax, as

shifted_sum_jit = jax.jit(shifted_sum, static_argnums=(2,))
jax.vmap(shifted_sum_jit, in_axes=(None, None, 0))(x, y, jnp.arange(1, S)).sum()

However the above does not work and gives me a ValueError: Non-hashable static arguments are not supported. What am I doing wrong? This seems a really trivial example that jax should be able to take care of.

I am using jax 0.2.24

[jakevdp]

Thanks for the question!

The reason for this error is that a parameter that is vmapped over is not static, so you cannot pass such a parameter to a function that is marked as requiring a static argument.

The fix is to not mark the argument static if you want to vmap over it.

The problem here, though, is that you're passing this argument (s) to another function that requires a static argument: index slicing. In JAX the size of an array must be static, so the s in x[:s] cannot be a traced value.

The solution here, if you want to sum over dynamically-sized slices, is to write your code in a way that does not require generating dynamically-sized arrays. Here is a workaround:

import jax
import jax.numpy as jnp

def shifted_sum(x, y, s):
  x_sum = jnp.where(jnp.arange(len(x)) < s, x, 0).sum()
  y_sum = jnp.where(jnp.arange(len(y)) >= len(y) - s, y, 0).sum()
  return x_sum + y_sum

x = y = jnp.arange(5)
S = 4

shifted_sum_jit = jax.jit(shifted_sum)
jax.vmap(shifted_sum_jit, in_axes=(None, None, 0))(x, y, jnp.arange(1, S)).sum()
# DeviceArray(24, dtype=int32)

There is some ongoing work toward supporting dynamically-shaped arrays within jit and other transforms; search "dynamic shapes" in JAX's issue tracker to see some of the progress. But for now, when you want to work with dynamically-shaped arrays, you have to use these kinds of workarounds.

[guglielmogattiglio]

Thanks for your answer! I see what you did there, however my function is actually more complex than that because the two slices of x and y interact in a way that I cannot apply the same workaround (if I took the dot product x[s:] @ y[:-s] instead of the sum, I would need to pad one of the two arrays to make the workaround work).

I understand that apparently the vmapped parameter is not static, not even if I use a static input np.arange(1, S), or if I wrap everything inside another function and mark S as static (which in my opinion should work since we know how long np.arange(1, S) is and how many calls to shifted_sum() should be made).

To sum up, I have two more questions:

1. Is there no generic strategy to use vmap like above but for a generic shifted_sum() function having a similar structure (behavior dependent on a static parameter)?
2. Is there no other strategy within JAX that I can use to call shifted_sum() in parallel for a range of s, and then collect the result? I am working on a single GPU, so would ideally like to spread the computation across the GPU cores (which is exactly what vmap does, as far as I understand).

Thanks a lot for your invaluable help!

[jakevdp]

vmap is implemented by tracing over mapped variables. So by design, they will be dynamic, not static. If you need your mapped-over values to be static, you cannot do so via any JAX transformation (a Python for-loop might be sufficient depending how many batches you have).

[jakevdp]

But to take advantage of GPU parallelism, I'd suggest exploring some padding-based option analogous to the sum example above. This is not just JAX advice, but advice about how to take advantage of the GPU hardware! Until JAX has better support for dynamic shapes (which will likely involve creating that sort of padded computation under the hood) then unfortunately the best approach is to construct such computations manually.

[guglielmogattiglio]

Ok perfect, this clarifies everything. Thank you so much!