def _GetArgsMaker(self, rng, shapes, dtypes, onp_arrays=True):

    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]
    return f
def _GetArgsMaker(self, rng, shapes, dtypes, onp_arrays=True):

    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]
    return f