def f():
    out = [rng(shape, (dtype or jnp.float_)) for (shape, dtype) in zip(shapes, dtypes)]
    if onp_arrays:
        return out
    return [(jnp.asarray(a) if isinstance(a, (onp.ndarray, onp.generic)) else a) for a in out]
def f():
    out = [rng(shape, (dtype or jnp.float_)) for (shape, dtype) in zip(shapes, dtypes)]
    if onp_arrays:
        return out
    return [(jnp.asarray(a) if isinstance(a, (onp.ndarray, onp.generic)) else a) for a in out]