Dropout layer doesn't work with jax backend when using jit + StatelessScope in training mode

[tirthasheshpatel]

The Dropout layer throws an error 'list' object has no attribute 'dtype' when run with jit and StatelessScope enabled. It looks like the seed's state value isn't getting initialized properly.

Here's an MRE:

import keras_core
from keras_core import backend
from keras_core.operations import numpy as knp
from jax import numpy as jnp
import jax

x = knp.array([0.1, 0.2, 0.3])

@jax.jit
def train_step(x):
    with backend.StatelessScope():
        x = keras_core.layers.Dropout(rate=0.1)(x, training=True)
    return x

keras_core.utils.traceback_utils.disable_traceback_filtering()
x = train_step(x)
assert isinstance(x, jnp.ndarray)