dask-array 0.1.0__py3-none-any.whl

dask_array/random/_expr.py

@@ -0,0 +1,256 @@
+from __future__ import annotations
+
+import importlib
+from itertools import product
+
+import numpy as np
+
+from dask._task_spec import Dict as TaskDict
+from dask._task_spec import Task, TaskRef
+from dask._task_spec import Tuple as TaskTuple
+from dask_array.io import IO
+from dask_array._core_utils import broadcast_shapes, normalize_chunks
+from dask_array._backends_array import array_creation_dispatch
+from dask.tokenize import tokenize
+from dask.utils import cached_property, random_state_data, typename
+
+from ._generator import Generator
+from ._random_state import RandomState
+
+
+def _rng_from_bitgen(bitgen):
+    # Assumes typename(bitgen) starts with importable
+    # library name (e.g. "numpy" or "cupy")
+    backend_name = typename(bitgen).split(".")[0]
+    backend_lib = importlib.import_module(backend_name)
+    return backend_lib.random.default_rng(bitgen)
+
+
+def _spawn_bitgens(bitgen, n_bitgens):
+    seeds = bitgen._seed_seq.spawn(n_bitgens)
+    bitgens = [type(bitgen)(seed) for seed in seeds]
+    return bitgens
+
+
+def _apply_random_func(rng, funcname, bitgen, size, args, kwargs):
+    """Apply random module method with seed"""
+    if isinstance(bitgen, np.random.SeedSequence):
+        bitgen = rng(bitgen)
+    rng = _rng_from_bitgen(bitgen)
+    func = getattr(rng, funcname)
+    return func(*args, size=size, **kwargs)
+
+
+def _apply_random(RandomState, funcname, state_data, size, args, kwargs):
+    """Apply RandomState method with seed"""
+    if RandomState is None:
+        RandomState = array_creation_dispatch.RandomState
+    state = RandomState(state_data)
+    func = getattr(state, funcname)
+    return func(*args, size=size, **kwargs)
+
+
+class Random(IO):
+    _parameters = [
+        "rng",
+        "distribution",
+        "size",
+        "chunks",
+        "extra_chunks",
+        "args",
+        "kwargs",
+    ]
+    _defaults = {"extra_chunks": ()}
+    _is_blockwise_fusable = True
+
+    @cached_property
+    def kwargs(self):
+        return self.operand("kwargs")
+
+    @cached_property
+    def _base_chunks(self):
+        """Chunks for the size dimensions, excluding extra_chunks."""
+        size = self.operand("size")
+        chunks = self.operand("chunks")
+        shape = broadcast_shapes(size) if size is not None else ()
+        return normalize_chunks(
+            chunks,
+            shape,
+            dtype=self.kwargs.get("dtype", np.float64),
+        )
+
+    @property
+    def chunks(self):
+        return self._base_chunks + self.extra_chunks
+
+    @cached_property
+    def _info(self):
+        sizes = list(product(*self._base_chunks))
+        if isinstance(self.rng, Generator):
+            bitgens = _spawn_bitgens(self.rng._bit_generator, len(sizes))
+            bitgen_token = tokenize(bitgens)
+            bitgens = [_bitgen._seed_seq for _bitgen in bitgens]
+            func_applier = _apply_random_func
+            gen = type(self.rng._bit_generator)
+        elif isinstance(self.rng, RandomState):
+            bitgens = random_state_data(len(sizes), self.rng._numpy_state)
+            bitgen_token = tokenize(bitgens)
+            func_applier = _apply_random
+            gen = self.rng._RandomState
+        else:
+            raise TypeError("Unknown object type: Not a Generator and Not a RandomState")
+        token = tokenize(bitgen_token, self.size, self.chunks, self.args, self.kwargs)
+        name = f"{self.distribution}-{token}"
+
+        return bitgens, name, sizes, gen, func_applier
+
+    @property
+    def _name(self):
+        return self._info[1]
+
+    @property
+    def bitgens(self):
+        return self._info[0]
+
+    def _layer(self):
+        result = {}
+        for block_id in product(*[range(len(bd)) for bd in self.chunks]):
+            key = (self._name, *block_id)
+            result[key] = self._task(key, block_id)
+        return result
+
+    def _block_id_to_flat_index(self, block_id: tuple[int, ...]) -> int:
+        """Convert N-D block_id to flat index for bitgens/sizes lookup."""
+        # Only use base_chunks dimensions (exclude extra_chunks which are always 0)
+        base_block_id = block_id[: len(self._base_chunks)]
+        flat_idx = 0
+        stride = 1
+        for i in reversed(range(len(base_block_id))):
+            flat_idx += base_block_id[i] * stride
+            stride *= len(self._base_chunks[i])
+        return flat_idx
+
+    def _task(self, key, block_id: tuple[int, ...]) -> Task:
+        """Generate task for a specific output block."""
+        from dask_array._expr import ArrayExpr
+
+        bitgens, name, sizes, gen, func_applier = self._info
+        flat_idx = self._block_id_to_flat_index(block_id)
+
+        # Convert array expressions to TaskRefs
+        task_args = []
+        has_array_args = False
+        for arg in self.args:
+            if isinstance(arg, ArrayExpr):
+                # Reference the corresponding block of the dependency
+                task_args.append(TaskRef((arg._name,) + block_id))
+                has_array_args = True
+            else:
+                task_args.append(arg)
+
+        task_kwargs = {}
+        has_array_kwargs = False
+        for k, v in self.kwargs.items():
+            if isinstance(v, ArrayExpr):
+                task_kwargs[k] = TaskRef((v._name,) + block_id)
+                has_array_kwargs = True
+            else:
+                task_kwargs[k] = v
+
+        # Use TaskTuple/TaskDict to properly track dependencies inside containers
+        args_container = TaskTuple(*task_args) if has_array_args else tuple(task_args)
+        kwargs_container = TaskDict(task_kwargs) if has_array_kwargs else task_kwargs
+
+        return Task(
+            key,
+            func_applier,
+            gen,
+            self.distribution,
+            bitgens[flat_idx],
+            sizes[flat_idx],
+            args_container,
+            kwargs_container,
+        )
+
+    def dependencies(self):
+        """Return array expression dependencies."""
+        from dask_array._expr import ArrayExpr
+
+        deps = []
+        for arg in self.args:
+            if isinstance(arg, ArrayExpr):
+                deps.append(arg)
+        for v in self.kwargs.values():
+            if isinstance(v, ArrayExpr):
+                deps.append(v)
+        return deps
+
+    def _input_block_id(self, dep, block_id: tuple[int, ...]) -> tuple[int, ...]:
+        """Map output block_id to input block_id for dependencies."""
+        return block_id
+
+    @cached_property
+    def _meta(self):
+        from dask_array._expr import ArrayExpr
+
+        bitgens, name, sizes, gen, func_applier = self._info
+        size = self.operand("size")
+        meta_size = (0,) * len(size) if size is not None else ()
+
+        # Convert array arguments to scalars for meta computation
+        def to_scalar(x):
+            if isinstance(x, ArrayExpr):
+                # Array expression - use dtype to create a scalar
+                return x.dtype.type(1)
+            elif isinstance(x, np.ndarray):
+                # Numpy array
+                return x.flat[0]
+            return x
+
+        meta_args = tuple(to_scalar(arg) for arg in self.args)
+        meta_kwargs = {k: to_scalar(v) for k, v in self.kwargs.items()}
+
+        return func_applier(
+            gen,
+            self.distribution,
+            bitgens[0],
+            meta_size,
+            meta_args,
+            meta_kwargs,
+        )
+
+
+# Distribution-specific subclasses with explicit parameters
+# These avoid storing array dependencies inside tuple operands
+
+
+class RandomNormal(Random):
+    """Normal distribution with explicit loc and scale parameters."""
+
+    _parameters = ["rng", "size", "chunks", "extra_chunks", "loc", "scale"]
+    _defaults = {"extra_chunks": (), "loc": 0.0, "scale": 1.0}
+    distribution = "normal"
+
+    @property
+    def args(self):
+        return (self.loc, self.scale)
+
+    @property
+    def kwargs(self):
+        return {}
+
+
+class RandomPoisson(Random):
+    """Poisson distribution with explicit lam parameter."""
+
+    _parameters = ["rng", "size", "chunks", "extra_chunks", "lam"]
+    _defaults = {"extra_chunks": (), "lam": 1.0}
+    distribution = "poisson"
+
+    @property
+    def args(self):
+        return (self.lam,)
+
+    @property
+    def kwargs(self):
+        return {}

dask_array/random/_generator.py

@@ -0,0 +1,441 @@
+from __future__ import annotations
+
+import importlib
+import numbers
+
+import numpy as np
+
+from dask_array._new_collection import new_collection
+from dask_array.creation import arange
+from dask_array._backends_array import array_creation_dispatch
+from dask.utils import derived_from, typename
+
+from ._utils import _wrap_func
+
+
+class Generator:
+    """
+    Container for the BitGenerators.
+
+    ``Generator`` exposes a number of methods for generating random
+    numbers drawn from a variety of probability distributions and serves
+    as a replacement for ``RandomState``. The main difference between the
+    two is that ``Generator`` relies on an additional ``BitGenerator`` to
+    manage state and generate the random bits, which are then transformed
+    into random values from useful distributions. The default ``BitGenerator``
+    used by ``Generator`` is ``PCG64``. The ``BitGenerator`` can be changed
+    by passing an instantiated ``BitGenerator`` to ``Generator``.
+
+    The function :func:`dask.array.random.default_rng` is the recommended way
+    to instantiate a ``Generator``.
+
+    .. warning::
+
+       No Compatibility Guarantee.
+
+       ``Generator`` does not provide a version compatibility guarantee. In
+       particular, as better algorithms evolve the bit stream may change.
+
+    Parameters
+    ----------
+    bit_generator : BitGenerator
+        BitGenerator to use as the core generator.
+
+    Notes
+    -----
+    In addition to the distribution-specific arguments, each ``Generator``
+    method takes a keyword argument `size` that defaults to ``None``. If
+    `size` is ``None``, then a single value is generated and returned. If
+    `size` is an integer, then a 1-D array filled with generated values is
+    returned. If `size` is a tuple, then an array with that shape is
+    filled and returned.
+
+    The Python stdlib module `random` contains pseudo-random number generator
+    with a number of methods that are similar to the ones available in
+    ``Generator``. It uses Mersenne Twister, and this bit generator can
+    be accessed using ``MT19937``. ``Generator``, besides being
+    Dask-aware, has the advantage that it provides a much larger number
+    of probability distributions to choose from.
+
+    All ``Generator`` methods are identical to ``np.random.Generator`` except
+    that they also take a `chunks=` keyword argument.
+
+    ``Generator`` does not guarantee parity in the generated numbers
+    with any third party library. In particular, numbers generated by
+    `Dask` and `NumPy` will differ even if they use the same seed.
+
+    Examples
+    --------
+    >>> from numpy.random import PCG64
+    >>> from dask_array.random import Generator
+    >>> rng = Generator(PCG64())
+    >>> rng.standard_normal().compute() # doctest: +SKIP
+    array(0.44595957)  # random
+
+    See Also
+    --------
+    default_rng : Recommended constructor for `Generator`.
+    np.random.Generator
+    """
+
+    def __init__(self, bit_generator):
+        self._bit_generator = bit_generator
+
+    def __str__(self):
+        return f"{self.__class__.__name__}({self._bit_generator.__class__.__name__})"
+
+    @property
+    def _backend_name(self):
+        # Assumes typename(self._RandomState) starts with an
+        # array-library name (e.g. "numpy" or "cupy")
+        return typename(self._bit_generator).split(".")[0]
+
+    @property
+    def _backend(self):
+        # Assumes `self._backend_name` is an importable
+        # array-library name (e.g. "numpy" or "cupy")
+        return importlib.import_module(self._backend_name)
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def beta(self, a, b, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "beta", a, b, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def binomial(self, n, p, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "binomial", n, p, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def chisquare(self, df, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "chisquare", df, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def choice(
+        self,
+        a,
+        size=None,
+        replace=True,
+        p=None,
+        axis=0,
+        shuffle=True,
+        chunks="auto",
+    ):
+        from ._choice import RandomChoiceGenerator, _choice_validate_params
+
+        (
+            a_val,
+            a_expr,
+            size,
+            replace,
+            p_expr,
+            axis,
+            chunks,
+            meta,
+        ) = _choice_validate_params(self, a, size, replace, p, axis, chunks)
+
+        return new_collection(
+            RandomChoiceGenerator(
+                a_val,
+                a_expr,
+                chunks,
+                meta,
+                self._bit_generator,
+                replace,
+                p_expr,
+                axis,
+                shuffle,
+            )
+        )
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def exponential(self, scale=1.0, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "exponential", scale, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def f(self, dfnum, dfden, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "f", dfnum, dfden, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def gamma(self, shape, scale=1.0, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "gamma", shape, scale, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def geometric(self, p, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "geometric", p, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def gumbel(self, loc=0.0, scale=1.0, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "gumbel", loc, scale, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def hypergeometric(self, ngood, nbad, nsample, size=None, chunks="auto", **kwargs):
+        return _wrap_func(
+            self,
+            "hypergeometric",
+            ngood,
+            nbad,
+            nsample,
+            size=size,
+            chunks=chunks,
+            **kwargs,
+        )
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def integers(
+        self,
+        low,
+        high=None,
+        size=None,
+        dtype=np.int64,
+        endpoint=False,
+        chunks="auto",
+        **kwargs,
+    ):
+        return _wrap_func(
+            self,
+            "integers",
+            low,
+            high=high,
+            size=size,
+            dtype=dtype,
+            endpoint=endpoint,
+            chunks=chunks,
+            **kwargs,
+        )
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def laplace(self, loc=0.0, scale=1.0, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "laplace", loc, scale, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def logistic(self, loc=0.0, scale=1.0, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "logistic", loc, scale, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def lognormal(self, mean=0.0, sigma=1.0, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "lognormal", mean, sigma, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def logseries(self, p, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "logseries", p, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def multinomial(self, n, pvals, size=None, chunks="auto", **kwargs):
+        return _wrap_func(
+            self,
+            "multinomial",
+            n,
+            pvals,
+            size=size,
+            chunks=chunks,
+            extra_chunks=((len(pvals),),),
+            **kwargs,
+        )
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def multivariate_hypergeometric(self, colors, nsample, size=None, method="marginals", chunks="auto", **kwargs):
+        return _wrap_func(
+            self,
+            "multivariate_hypergeometric",
+            colors,
+            nsample,
+            size=size,
+            method=method,
+            chunks=chunks,
+            **kwargs,
+        )
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def negative_binomial(self, n, p, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "negative_binomial", n, p, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def noncentral_chisquare(self, df, nonc, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "noncentral_chisquare", df, nonc, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def noncentral_f(self, dfnum, dfden, nonc, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "noncentral_f", dfnum, dfden, nonc, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def normal(self, loc=0.0, scale=1.0, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "normal", loc, scale, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def pareto(self, a, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "pareto", a, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def permutation(self, x):
+        from dask_array.slicing._utils import shuffle_slice
+
+        from ._utils import _shuffle
+
+        if self._backend_name == "cupy":
+            raise NotImplementedError(
+                "`Generator.permutation` not supported for cupy-backed "
+                "Generator objects. Use the 'numpy' array backend to "
+                "call `dask.array.random.default_rng`, or pass in "
+                " `numpy.random.PCG64()`."
+            )
+
+        if isinstance(x, numbers.Number):
+            x = arange(x, chunks="auto")
+
+        index = self._backend.arange(len(x))
+        _shuffle(self._bit_generator, index)
+        return shuffle_slice(x, index)
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def poisson(self, lam=1.0, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "poisson", lam, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def power(self, a, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "power", a, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def random(self, size=None, dtype=np.float64, out=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "random", size=size, dtype=dtype, out=out, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def rayleigh(self, scale=1.0, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "rayleigh", scale, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def standard_cauchy(self, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "standard_cauchy", size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def standard_exponential(self, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "standard_exponential", size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def standard_gamma(self, shape, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "standard_gamma", shape, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def standard_normal(self, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "standard_normal", size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def standard_t(self, df, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "standard_t", df, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def triangular(self, left, mode, right, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "triangular", left, mode, right, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def uniform(self, low=0.0, high=1.0, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "uniform", low, high, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def vonmises(self, mu, kappa, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "vonmises", mu, kappa, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def wald(self, mean, scale, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "wald", mean, scale, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def weibull(self, a, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "weibull", a, size=size, chunks=chunks, **kwargs)
+
+    @derived_from(np.random.Generator, skipblocks=1)
+    def zipf(self, a, size=None, chunks="auto", **kwargs):
+        return _wrap_func(self, "zipf", a, size=size, chunks=chunks, **kwargs)
+
+
+def default_rng(seed=None):
+    """
+    Construct a new Generator with the default BitGenerator (PCG64).
+
+    Parameters
+    ----------
+    seed : {None, int, array_like[ints], SeedSequence, BitGenerator, Generator}, optional
+        A seed to initialize the `BitGenerator`. If None, then fresh,
+        unpredictable entropy will be pulled from the OS. If an ``int`` or
+        ``array_like[ints]`` is passed, then it will be passed to
+        `SeedSequence` to derive the initial `BitGenerator` state. One may
+        also pass in a `SeedSequence` instance.
+        Additionally, when passed a `BitGenerator`, it will be wrapped by
+        `Generator`. If passed a `Generator`, it will be returned unaltered.
+
+    Returns
+    -------
+    Generator
+        The initialized generator object.
+
+    Notes
+    -----
+    If ``seed`` is not a `BitGenerator` or a `Generator`, a new
+    `BitGenerator` is instantiated. This function does not manage a default
+    global instance.
+
+    Examples
+    --------
+    ``default_rng`` is the recommended constructor for the random number
+    class ``Generator``. Here are several ways we can construct a random
+    number generator using ``default_rng`` and the ``Generator`` class.
+
+    Here we use ``default_rng`` to generate a random float:
+
+    >>> import dask_array as da
+    >>> rng = da.random.default_rng(12345)
+    >>> print(rng)
+    Generator(PCG64)
+    >>> rfloat = rng.random().compute()
+    >>> rfloat
+    array(0.86999885)
+    >>> type(rfloat)
+    <class 'numpy.ndarray'>
+
+    Here we use ``default_rng`` to generate 3 random integers between 0
+    (inclusive) and 10 (exclusive):
+
+    >>> import dask_array as da
+    >>> rng = da.random.default_rng(12345)
+    >>> rints = rng.integers(low=0, high=10, size=3).compute()
+    >>> rints
+    array([2, 8, 7])
+    >>> type(rints[0])
+    <class 'numpy.int64'>
+
+    Here we specify a seed so that we have reproducible results:
+
+    >>> import dask_array as da
+    >>> rng = da.random.default_rng(seed=42)
+    >>> print(rng)
+    Generator(PCG64)
+    >>> arr1 = rng.random((3, 3)).compute()
+    >>> arr1
+    array([[0.91674416, 0.91098667, 0.8765925 ],
+           [0.30931841, 0.95465607, 0.17509458],
+           [0.99662814, 0.75203348, 0.15038118]])
+
+    If we exit and restart our Python interpreter, we'll see that we
+    generate the same random numbers again:
+
+    >>> import dask_array as da
+    >>> rng = da.random.default_rng(seed=42)
+    >>> arr2 = rng.random((3, 3)).compute()
+    >>> arr2
+    array([[0.91674416, 0.91098667, 0.8765925 ],
+           [0.30931841, 0.95465607, 0.17509458],
+           [0.99662814, 0.75203348, 0.15038118]])
+
+    See Also
+    --------
+    np.random.default_rng
+    """
+    if hasattr(seed, "capsule"):
+        # We are passed a BitGenerator, so just wrap it
+        return Generator(seed)
+    elif isinstance(seed, Generator):
+        # Pass through a Generator
+        return seed
+    elif hasattr(seed, "bit_generator"):
+        # a Generator. Just not ours
+        return Generator(seed.bit_generator)
+    # Otherwise, use the backend-default BitGenerator
+    return Generator(array_creation_dispatch.default_bit_generator(seed))