Can we use TrainState with pad_shard_upad?

[zjoukhadar]

I have a train_step that has TrainState as one of its argument
def train_step(state: train_state.TrainState,metric,batch): #some code return state, metrics

I am using pad_shard_unpad like this:

train_step = jax_utils.pad_shard_unpad(train_step, static_argnums=(0,1), static_return=True)

My state:TrainState has optax optimizer in it.

When I run my code. I am getting the following error:

Non-hashable static arguments are not supported. An error occurred during a call to 'train_step' while trying to hash an object of type <class 'jaxlib.xla_extension.DeviceArray'>, [ 0 42]. The error was: TypeError: unhashable type: 'DeviceArray'

[marcvanzee]

Are you using any other JAX transformations in your code as well? This error happens when you are passing a JAX array as an argument to a jitted function, but you have specified that this argument is static. JAX arrays should not be passed as compile-time constants since they cannot be hashed.

Here's a minimal example explaining what I mean:

import jax

jax.jit(lambda x: x + 1, static_argnums=(0,))(jax.numpy.array([0, 1]))  # same error

jax.jit(lambda x: x + 1)(jax.numpy.array([0, 1]))  # works

[zjoukhadar]

Thanks @marcvanzee
I think I had a bug in my code. I was mistakenly passing rng which is jax.random.PRNGKey(seed) which is a devicearray and I was asking it to be static. devicearray cannot be passed as static arguments. I will close this issue now.