How can I make a KeyArray hashable under the new system?

[NeilGirdhar]

I have some code:

        def hashable(x: Any) -> Hashable:
            if is_dataclass(x):
                return astuple(x)
            if isinstance(x, KeyArray):
                return tuple(x.unsafe_raw_array().tolist())
            ...

Is there a way to make this work under Jax's new KeyArray that acts like an ordinary array?

[jakevdp]

Rather than isinstance(x, KeyArray), you can use issubdtype(x.dtype, dtypes.prng_key). See #17297 for some discussion.

[jakevdp]

Also, use jax.random.key_data rather than unsafe_raw_array

[NeilGirdhar]

Awesome, thanks for the quick reply as usual!