Destructing DeviceArray produces Numpy arrays

Don’t know if this is a new feature or a bug. Destructing a DeviceArray is resulting in its pieces being converted to Numpy arrays:

import jax
import jax.numpy as jnp

key = jax.random.PRNGKey(0)
key1, key2 = jax.random.split(key)

assert not isintance(key1, jnp.ndarray)
assert not isintance(key2, jnp.ndarray)

This is surprising given this pattern with split is used extensively in the docs. Now it seems you have to do this:

key = jax.random.PRNGKey(0)
splits = jax.random.split(key)
key1, key2 = splits[0], splits[1]

assert isintance(key1, jnp.ndarray)
assert isintance(key2, jnp.ndarray)

Which is less nice.

Versions:

name = "jax"
version = "0.2.21"

name = "jaxlib"
version = "0.1.71"

[jakevdp]

Yeah, the issue is much larger than split, unfortunately. We really need to get #3821 or something like it in, but it has been challenging for various reasons.

[cgarciae]

Alternatively, you can also replace jax.random.split with this:

def iter_split(key: jnp.ndarray, num: int = 2) -> tp.Tuple[jnp.ndarray, ...]:
    splits = jax.random.split(key, num)
    return tuple(splits[i] for i in range(num))