Error: Non-hashable static arguments for fori_loop

[sokrypton]

Hello,
I'm trying to use a fori_loop and the iterator in my function.
When I do this, I get the following error:
"'DynamicJaxprTracer' object cannot be interpreted as an integer"

When I tried setting: static_argnums=(0,), the error is now:
"Non-hashable static arguments are not supported"

I feel like I'm missing something any ideas?

Here is a google colab notebook with working example (using python for-loop) and error when I try fori_loop:
https://colab.research.google.com/drive/15DN56p0R1t_Env3ETvz0dhXEOhW-tu28

[jakevdp]

The issue is that the third argument to jnp.eye must be static, and the arguments produced by fori_loop are by design non-static (i.e. traced). I would suggest finding another way to express your computation, for example via broadcasted indices. For example, this produces the same output:

import jax.numpy as jnp
i = jnp.arange(10)[:, None]
j = jnp.arange(5)
x = jnp.ones((10, 10))
x = x.at[i, i + j].set(2)

It is compatible with jit, and unlike fori_loop will have good performance on accelerators.

[sokrypton]

Thanks Jake! My real example is a bit more complicated, I just tried to think of an easy example to reproduce the error :D

I'll definitely look at ".at[]" and ".set()" was not aware of these!

For debugging the rest of my code, Is there an easy way to figure out which arguments in which function must be static?

[jakevdp]

Well, if you get a "'DynamicJaxprTracer' object cannot be interpreted as an integer" error that's a pretty good sign 😁 But I don't think functions are uniformly documented regarding which arguments must be static.

[sokrypton]

I'm stuck! :(
Seems nearly every function I use has some kind of static requirement.

Here is the actual problem I'm trying to solve. Any pointers would be great :D
https://colab.research.google.com/drive/1SsnOFyta1LH1xoByMCL9Tb3T57nwXS-d
(just a heads up my function is a little dumb, and currently only works when x.shape[0] <= x.shape[1] and both sides are even)

What I'm trying to do is take a matrix, and walk through it one corner to the next, one diagonal element at a time. Resulting in a rotated matrix:

I was thinking to use lax.fori_loop to do this along with jnp.diagonal() function... but unfortunately, the offset of diagonal() is required to be static. I was thinking to replace jnp.diagonal() with some jnp.eye(k=offset) masking, but that also failed due to k being static... (as you pointed out in the original reply).

[jakevdp]

I think you could do something like this, if I'm understanding the goal correctly:

import jax.numpy as jnp

def rotate(x):
  assert x.ndim == 2
  output = jnp.zeros((sum(x.shape) // 2, sum(x.shape) - 1))
  i = jnp.arange(x.shape[0])[:, None]
  j = jnp.arange(x.shape[1])
  return output.at[(i + j) // 2, j - i + len(i) - 1].set(x)

x = jnp.arange(10*20).reshape(10,20)
x_rotated = rotate(x)

[sokrypton]

Wow! That's exactly what I need. Thanks jake! :D