DeviceArray as static argument raise non-hashable error

[yiminghwang]

Hi, there,

I feed the DeviceArray to the function with jit, the simplified one shows as follows and raises the following error:

from jax import jit
import jax.numpy as jnp
from functools import partial

@partial(jit, static_argnums=(1,))
def f(x, v):
    return x@v

def main():
    A = jnp.ones((5, 5))
    x = jnp.array([1, 2, 3, 4, 5])
    y = f(A, x)
    print("end")

if __name__ == '__main__':
    main()

ValueError: Non-hashable static arguments are not supported. An error occurred during a call to 'f' while trying to hash an object of type <class 'jaxlib.xla_extension.DeviceArray'>, [1 2 3 4 5]. The error was:
TypeError: unhashable type: 'DeviceArray'

Does anyone know how to fix it?

[mattjj]

Thanks for the question!

Take a look at this comment. The thread (especially just above there) has an explanation for why this error is raised.

[yiminghwang]

Thanks for the quick response!