Matrix Multiplication is not running with Pallas

[bsaoptima]

Hi everyone,

So I had this Jax code which worked well and I am trying to use Pallas and I can't understand where my mistake is coming from.

key = jax.random.PRNGKey(50)
other = jax.random.normal(key, (10, 10))
diags = jax.random.normal(key, (3, 10))
offsets = (-2, 1, 2)

@functools.partial(jax.jit, static_argnums=(1,))
def dia_matmul(diags: Array, offsets: tuple[int], other:Array) -> Array:    
    return pl.pallas_call(
        dia_matmul_kernel,
        out_shape=jax.ShapeDtypeStruct(other.shape, other.dtype)
    )(diags, offsets, other)
dia_matmul(diags, offsets, other)

def dia_matmul_kernel(diags_ref, offsets_ref, other_ref, o_ref):
    diags, offsets, other = diags_ref[...], offsets_ref[:], other_ref[...]
    N = other.shape[0]
    out = jnp.zeros((N, N))
    for offset, diag in zip(offsets, diags):
        start = jax.numpy.maximum(0, offset)
        end = jax.numpy.minimum(N, N + offset)
        out += jnp.diag(diag[start:end], k=offset)
    o_ref[...] = out

The error is coming from the start = jax.numpy.maximum(0, offset) line: AssertionError: (ShapedArray(int32[]), MemRef<None>{int32[]})

I am not yet an expert at Jax and how the whole Tracing system works but the static_argnums should protect my offset tuple? It worked fine in simple Jax but now trying to do it in Pallas gives me those errors.

If anyone has an idea, it would be very helpful! Many thanks!