Change key type from prng_key to Array

[fehiepsi]

A key has type Array while key.dtype is a subclass of jax.dtypes.prng_key.

[juanitorduz]

LGTM :)

(can you please ensure the lint is happy :) )

[Qazalbash]

LGTM!

Can we use an alias like PRNGKeyArray to distinguish between regular arrays and keys?

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[fehiepsi]

I'm not sure if PRNGKey is an alias though. It seems jax.random.key returns an Array: https://github.com/jax-ml/jax/blob/2df8eced12c0e5a2456ebc0c693077109fa51fbd/jax/_src/random.py#L212-L214

[Qazalbash]

I'm not sure if PRNGKey is an alias though. It seems jax.random.key returns an Array: jax-ml/jax@2df8ece/jax/_src/random.py#L212-L214

I am not saying the alisa exists. My point was to use,

PRNGKeyArray: TypeAlias = Array

and use PRNGKeyArray when we are referring to Arrays that are keys.

[fehiepsi]

Gotcha. I'm fine with either, slightly prefer using Array to be aligned with jax and shorter - the variable name already convey the "key" message.

[Qazalbash]

Let's see this in the future :)