brainstate 0.0.2.post20240913__py2.py3-none-any.whl → 0.0.2.post20241010__py2.py3-none-any.whl

brainstate/{random.py → random/_rand_funs.py}

@@ -16,31 +16,14 @@
 
 # -*- coding: utf-8 -*-
 
-from collections import namedtuple
-from contextlib import contextmanager
-from functools import partial
-from operator import index
 from typing import Optional
 
-import brainunit as bu
-import jax
-import jax.numpy as jnp
-import jax.random as jr
 import numpy as np
-from jax import jit, vmap
-from jax import lax, core, dtypes
 
-from brainstate import environ
-from ._random_for_unit import uniform_for_unit, permutation_for_unit
-from ._state import State
-from .transform._error_if import jit_error_if
-from .typing import DTypeLike, Size, SeedOrKey
+from brainstate.typing import DTypeLike, Size, SeedOrKey
+from ._rand_state import RandomState, DEFAULT
 
 __all__ = [
-  'RandomState', 'DEFAULT',
-
-  'seed', 'default_rng', 'split_key', 'split_keys', 'seed_context',
-
   # numpy compatibility
   'rand', 'randint', 'random_integers', 'randn', 'random',
   'random_sample', 'ranf', 'sample', 'choice', 'permutation', 'shuffle', 'beta',
@@ -58,1186 +41,6 @@ __all__ = [
 ]
 
 
-class RandomState(State):
-  """RandomState that track the random generator state. """
-  __slots__ = ()
-
-  def __init__(self, seed_or_key: Optional[SeedOrKey] = None):
-    """RandomState constructor.
-
-    Parameters
-    ----------
-    seed_or_key: int, Array, optional
-      It can be an integer for initial seed of the random number generator,
-      or it can be a JAX's PRNKey, which is an array with two elements and `uint32` dtype.
-    """
-    with jax.ensure_compile_time_eval():
-      if seed_or_key is None:
-        seed_or_key = np.random.randint(0, 100000, 2, dtype=np.uint32)
-    if isinstance(seed_or_key, int):
-      key = jr.PRNGKey(seed_or_key)
-    else:
-      if len(seed_or_key) != 2 and seed_or_key.dtype != np.uint32:
-        raise ValueError('key must be an array with dtype uint32. '
-                         f'But we got {seed_or_key}')
-      key = seed_or_key
-    super().__init__(key)
-
-  def __repr__(self) -> str:
-    print_code = repr(self.value)
-    i = print_code.index('(')
-    return f'{self.__class__.__name__}(key={print_code[i:]})'
-
-  def _check_if_deleted(self):
-    if isinstance(self._value, jax.Array) and not isinstance(self._value, jax.core.Tracer) and self._value.is_deleted():
-      self.seed()
-
-  # ------------------- #
-  # seed and random key #
-  # ------------------- #
-
-  def clone(self):
-    return type(self)(self.split_key())
-
-  def seed(self, seed_or_key: Optional[SeedOrKey] = None):
-    """Sets a new random seed.
-
-    Parameters
-    ----------
-    seed_or_key: int, ArrayLike, optional
-      It can be an integer for initial seed of the random number generator,
-      or it can be a JAX's PRNKey, which is an array with two elements and `uint32` dtype.
-    """
-    if seed_or_key is None:
-      seed_or_key = np.random.randint(0, 100000, 2, dtype=np.uint32)
-    if isinstance(seed_or_key, int):
-      key = jr.PRNGKey(seed_or_key)
-    else:
-      if len(seed_or_key) != 2 and seed_or_key.dtype != np.uint32:
-        raise ValueError('key must be an array with dtype uint32. '
-                         f'But we got {seed_or_key}')
-      key = seed_or_key
-    self.value = key
-
-  def split_key(self):
-    """Create a new seed from the current seed.
-    """
-    if not isinstance(self.value, jax.Array):
-      self.value = jnp.asarray(self.value, dtype=jnp.uint32)
-    keys = jr.split(self.value, num=2)
-    self.value = keys[0]
-    return keys[1]
-
-  def split_keys(self, n: int):
-    """Create multiple seeds from the current seed. This is used
-    internally by `pmap` and `vmap` to ensure that random numbers
-    are different in parallel threads.
-
-    Parameters
-    ----------
-    n : int
-      The number of seeds to generate.
-    """
-    keys = jr.split(self.value, n + 1)
-    self.value = keys[0]
-    return keys[1:]
-
-  # ---------------- #
-  # random functions #
-  # ---------------- #
-
-  def rand(self, *dn, key: Optional[SeedOrKey] = None, dtype: DTypeLike = None):
-    key = self.split_key() if key is None else _formalize_key(key)
-    dtype = dtype or environ.dftype()
-    r = uniform_for_unit(key, shape=dn, minval=0., maxval=1., dtype=dtype)
-    return r
-
-  def randint(
-      self,
-      low,
-      high=None,
-      size: Optional[Size] = None,
-      dtype: DTypeLike = None,
-      key: Optional[SeedOrKey] = None
-  ):
-    if high is None:
-      high = low
-      low = 0
-    high = _check_py_seq(high)
-    low = _check_py_seq(low)
-    if size is None:
-      size = lax.broadcast_shapes(jnp.shape(low),
-                                  jnp.shape(high))
-    key = self.split_key() if key is None else _formalize_key(key)
-    dtype = dtype or environ.ditype()
-    r = jr.randint(key,
-                   shape=_size2shape(size),
-                   minval=low, maxval=high, dtype=dtype)
-    return r
-
-  def random_integers(
-      self,
-      low,
-      high=None,
-      size: Optional[Size] = None,
-      key: Optional[SeedOrKey] = None,
-      dtype: DTypeLike = None,
-  ):
-    low = _check_py_seq(low)
-    high = _check_py_seq(high)
-    if high is None:
-      high = low
-      low = 1
-    high += 1
-    if size is None:
-      size = lax.broadcast_shapes(jnp.shape(low), jnp.shape(high))
-    key = self.split_key() if key is None else _formalize_key(key)
-    dtype = dtype or environ.ditype()
-    r = jr.randint(key,
-                   shape=_size2shape(size),
-                   minval=low,
-                   maxval=high,
-                   dtype=dtype)
-    return r
-
-  def randn(self, *dn, key: Optional[SeedOrKey] = None, dtype: DTypeLike = None):
-    key = self.split_key() if key is None else _formalize_key(key)
-    dtype = dtype or environ.dftype()
-    r = jr.normal(key, shape=dn, dtype=dtype)
-    return r
-
-  def random(self,
-             size: Optional[Size] = None,
-             key: Optional[SeedOrKey] = None,
-             dtype: DTypeLike = None):
-    dtype = dtype or environ.dftype()
-    key = self.split_key() if key is None else _formalize_key(key)
-    r = uniform_for_unit(key, shape=_size2shape(size), minval=0., maxval=1., dtype=dtype)
-    return r
-
-  def random_sample(self,
-                    size: Optional[Size] = None,
-                    key: Optional[SeedOrKey] = None,
-                    dtype: DTypeLike = None):
-    r = self.random(size=size, key=key, dtype=dtype)
-    return r
-
-  def ranf(self,
-           size: Optional[Size] = None,
-           key: Optional[SeedOrKey] = None,
-           dtype: DTypeLike = None):
-    r = self.random(size=size, key=key, dtype=dtype)
-    return r
-
-  def sample(self,
-             size: Optional[Size] = None,
-             key: Optional[SeedOrKey] = None,
-             dtype: DTypeLike = None):
-    r = self.random(size=size, key=key, dtype=dtype)
-    return r
-
-  def choice(self,
-             a,
-             size: Optional[Size] = None,
-             replace=True,
-             p=None,
-             key: Optional[SeedOrKey] = None):
-    a = _check_py_seq(a)
-    p = _check_py_seq(p)
-    key = self.split_key() if key is None else _formalize_key(key)
-    r = jr.choice(key, a=a, shape=_size2shape(size), replace=replace, p=p)
-    return r
-
-  def permutation(self,
-                  x,
-                  axis: int = 0,
-                  independent: bool = False,
-                  key: Optional[SeedOrKey] = None):
-    x = _check_py_seq(x)
-    key = self.split_key() if key is None else _formalize_key(key)
-    r = permutation_for_unit(key, x, axis=axis, independent=independent)
-    return r
-
-  def shuffle(self,
-              x,
-              axis=0,
-              key: Optional[SeedOrKey] = None):
-    key = self.split_key() if key is None else _formalize_key(key)
-    x = permutation_for_unit(key, x, axis=axis)
-    return x
-
-  def beta(self,
-           a,
-           b,
-           size: Optional[Size] = None,
-           key: Optional[SeedOrKey] = None,
-           dtype: DTypeLike = None):
-    a = _check_py_seq(a)
-    b = _check_py_seq(b)
-    if size is None:
-      size = lax.broadcast_shapes(jnp.shape(a), jnp.shape(b))
-    key = self.split_key() if key is None else _formalize_key(key)
-    dtype = dtype or environ.dftype()
-    r = jr.beta(key, a=a, b=b, shape=_size2shape(size), dtype=dtype)
-    return r
-
-  def exponential(self,
-                  scale=None,
-                  size: Optional[Size] = None,
-                  key: Optional[SeedOrKey] = None,
-                  dtype: DTypeLike = None):
-    if size is None:
-      size = jnp.shape(scale)
-    key = self.split_key() if key is None else _formalize_key(key)
-    dtype = dtype or environ.dftype()
-    scale = jnp.asarray(scale, dtype=dtype)
-    r = jr.exponential(key, shape=_size2shape(size), dtype=dtype)
-    if scale is not None:
-      r = r / scale
-    return r
-
-  def gamma(self,
-            shape,
-            scale=None,
-            size: Optional[Size] = None,
-            key: Optional[SeedOrKey] = None,
-            dtype: DTypeLike = None):
-    shape = _check_py_seq(shape)
-    scale = _check_py_seq(scale)
-    if size is None:
-      size = lax.broadcast_shapes(jnp.shape(shape), jnp.shape(scale))
-    key = self.split_key() if key is None else _formalize_key(key)
-    dtype = dtype or environ.dftype()
-    r = jr.gamma(key, a=shape, shape=_size2shape(size), dtype=dtype)
-    if scale is not None:
-      r = r * scale
-    return r
-
-  def gumbel(self,
-             loc=None,
-             scale=None,
-             size: Optional[Size] = None,
-             key: Optional[SeedOrKey] = None,
-             dtype: DTypeLike = None):
-    loc = _check_py_seq(loc)
-    scale = _check_py_seq(scale)
-    if size is None:
-      size = lax.broadcast_shapes(jnp.shape(loc), jnp.shape(scale))
-    key = self.split_key() if key is None else _formalize_key(key)
-    dtype = dtype or environ.dftype()
-    r = _loc_scale(loc, scale, jr.gumbel(key, shape=_size2shape(size), dtype=dtype))
-    return r
-
-  def laplace(self,
-              loc=None,
-              scale=None,
-              size: Optional[Size] = None,
-              key: Optional[SeedOrKey] = None,
-              dtype: DTypeLike = None):
-    loc = _check_py_seq(loc)
-    scale = _check_py_seq(scale)
-    if size is None:
-      size = lax.broadcast_shapes(jnp.shape(loc), jnp.shape(scale))
-    key = self.split_key() if key is None else _formalize_key(key)
-    dtype = dtype or environ.dftype()
-    r = _loc_scale(loc, scale, jr.laplace(key, shape=_size2shape(size), dtype=dtype))
-    return r
-
-  def logistic(self,
-               loc=None,
-               scale=None,
-               size: Optional[Size] = None,
-               key: Optional[SeedOrKey] = None,
-               dtype: DTypeLike = None):
-    loc = _check_py_seq(loc)
-    scale = _check_py_seq(scale)
-    if size is None:
-      size = lax.broadcast_shapes(jnp.shape(loc), jnp.shape(scale))
-    key = self.split_key() if key is None else _formalize_key(key)
-    dtype = dtype or environ.dftype()
-    r = _loc_scale(loc, scale, jr.logistic(key, shape=_size2shape(size), dtype=dtype))
-    return r
-
-  def normal(self,
-             loc=None,
-             scale=None,
-             size: Optional[Size] = None,
-             key: Optional[SeedOrKey] = None,
-             dtype: DTypeLike = None):
-    loc = _check_py_seq(loc)
-    scale = _check_py_seq(scale)
-    if size is None:
-      size = lax.broadcast_shapes(jnp.shape(scale), jnp.shape(loc))
-    key = self.split_key() if key is None else _formalize_key(key)
-    dtype = dtype or environ.dftype()
-    r = _loc_scale(loc, scale, jr.normal(key, shape=_size2shape(size), dtype=dtype))
-    return r
-
-  def pareto(self,
-             a,
-             size: Optional[Size] = None,
-             key: Optional[SeedOrKey] = None,
-             dtype: DTypeLike = None):
-    if size is None:
-      size = jnp.shape(a)
-    key = self.split_key() if key is None else _formalize_key(key)
-    dtype = dtype or environ.dftype()
-    a = jnp.asarray(a, dtype=dtype)
-    r = jr.pareto(key, b=a, shape=_size2shape(size), dtype=dtype)
-    return r
-
-  def poisson(self,
-              lam=1.0,
-              size: Optional[Size] = None,
-              key: Optional[SeedOrKey] = None,
-              dtype: DTypeLike = None):
-    lam = _check_py_seq(lam)
-    if size is None:
-      size = jnp.shape(lam)
-    key = self.split_key() if key is None else _formalize_key(key)
-    dtype = dtype or environ.ditype()
-    r = jr.poisson(key, lam=lam, shape=_size2shape(size), dtype=dtype)
-    return r
-
-  def standard_cauchy(self,
-                      size: Optional[Size] = None,
-                      key: Optional[SeedOrKey] = None,
-                      dtype: DTypeLike = None):
-    key = self.split_key() if key is None else _formalize_key(key)
-    dtype = dtype or environ.dftype()
-    r = jr.cauchy(key, shape=_size2shape(size), dtype=dtype)
-    return r
-
-  def standard_exponential(self,
-                           size: Optional[Size] = None,
-                           key: Optional[SeedOrKey] = None,
-                           dtype: DTypeLike = None):
-    key = self.split_key() if key is None else _formalize_key(key)
-    dtype = dtype or environ.dftype()
-    r = jr.exponential(key, shape=_size2shape(size), dtype=dtype)
-    return r
-
-  def standard_gamma(self,
-                     shape,
-                     size: Optional[Size] = None,
-                     key: Optional[SeedOrKey] = None,
-                     dtype: DTypeLike = None):
-    shape = _check_py_seq(shape)
-    if size is None:
-      size = jnp.shape(shape)
-    key = self.split_key() if key is None else _formalize_key(key)
-    dtype = dtype or environ.dftype()
-    r = jr.gamma(key, a=shape, shape=_size2shape(size), dtype=dtype)
-    return r
-
-  def standard_normal(self,
-                      size: Optional[Size] = None,
-                      key: Optional[SeedOrKey] = None,
-                      dtype: DTypeLike = None):
-    key = self.split_key() if key is None else _formalize_key(key)
-    dtype = dtype or environ.dftype()
-    r = jr.normal(key, shape=_size2shape(size), dtype=dtype)
-    return r
-
-  def standard_t(self, df,
-                 size: Optional[Size] = None,
-                 key: Optional[SeedOrKey] = None,
-                 dtype: DTypeLike = None):
-    df = _check_py_seq(df)
-    if size is None:
-      size = jnp.shape(size)
-    key = self.split_key() if key is None else _formalize_key(key)
-    dtype = dtype or environ.dftype()
-    r = jr.t(key, df=df, shape=_size2shape(size), dtype=dtype)
-    return r
-
-  def uniform(self,
-              low=0.0,
-              high=1.0,
-              size: Optional[Size] = None,
-              key: Optional[SeedOrKey] = None,
-              dtype: DTypeLike = None):
-    low = _check_py_seq(low)
-    high = _check_py_seq(high)
-    if size is None:
-      size = lax.broadcast_shapes(jnp.shape(low), jnp.shape(high))
-    key = self.split_key() if key is None else _formalize_key(key)
-    dtype = dtype or environ.dftype()
-    r = uniform_for_unit(key, shape=_size2shape(size), minval=low, maxval=high, dtype=dtype)
-    return r
-
-  def __norm_cdf(self, x, sqrt2, dtype):
-    # Computes standard normal cumulative distribution function
-    return (np.asarray(1., dtype) + lax.erf(x / sqrt2)) / np.asarray(2., dtype)
-
-  def truncated_normal(
-      self,
-      lower,
-      upper,
-      size: Optional[Size] = None,
-      loc=0.,
-      scale=1.,
-      key: Optional[SeedOrKey] = None,
-      dtype: DTypeLike = None
-  ):
-    lower = _check_py_seq(lower)
-    upper = _check_py_seq(upper)
-    loc = _check_py_seq(loc)
-    scale = _check_py_seq(scale)
-    dtype = dtype or environ.dftype()
-
-    lower = bu.math.asarray(lower, dtype=dtype)
-    upper = bu.math.asarray(upper, dtype=dtype)
-    loc = bu.math.asarray(loc, dtype=dtype)
-    scale = bu.math.asarray(scale, dtype=dtype)
-    unit = bu.get_unit(lower)
-    lower, upper, loc, scale = (
-      lower.mantissa if isinstance(lower, bu.Quantity) else lower,
-      bu.Quantity(upper).in_unit(unit).mantissa,
-      bu.Quantity(loc).in_unit(unit).mantissa,
-      bu.Quantity(scale).in_unit(unit).mantissa
-    )
-
-    jit_error_if(
-      bu.math.any(bu.math.logical_or(loc < lower - 2 * scale, loc > upper + 2 * scale)),
-      "mean is more than 2 std from [lower, upper] in truncated_normal. "
-      "The distribution of values may be incorrect."
-    )
-
-    if size is None:
-      size = bu.math.broadcast_shapes(jnp.shape(lower),
-                                      jnp.shape(upper),
-                                      jnp.shape(loc),
-                                      jnp.shape(scale))
-
-    # Values are generated by using a truncated uniform distribution and
-    # then using the inverse CDF for the normal distribution.
-    # Get upper and lower cdf values
-    sqrt2 = np.array(np.sqrt(2), dtype=dtype)
-    l = self.__norm_cdf((lower - loc) / scale, sqrt2, dtype)
-    u = self.__norm_cdf((upper - loc) / scale, sqrt2, dtype)
-
-    # Uniformly fill tensor with values from [l, u], then translate to
-    # [2l-1, 2u-1].
-    key = self.split_key() if key is None else _formalize_key(key)
-    out = uniform_for_unit(
-      key, size, dtype,
-      minval=lax.nextafter(2 * l - 1, np.array(np.inf, dtype=dtype)),
-      maxval=lax.nextafter(2 * u - 1, np.array(-np.inf, dtype=dtype))
-    )
-
-    # Use inverse cdf transform for normal distribution to get truncated
-    # standard normal
-    out = lax.erf_inv(out)
-
-    # Transform to proper mean, std
-    out = out * scale * sqrt2 + loc
-
-    # Clamp to ensure it's in the proper range
-    out = jnp.clip(
-      out,
-      lax.nextafter(lax.stop_gradient(lower), np.array(np.inf, dtype=dtype)),
-      lax.nextafter(lax.stop_gradient(upper), np.array(-np.inf, dtype=dtype))
-    )
-    return out if unit.is_unitless else bu.Quantity(out, unit=unit)
-
-  def _check_p(self, p):
-    raise ValueError(f'Parameter p should be within [0, 1], but we got {p}')
-
-  def bernoulli(self,
-                p,
-                size: Optional[Size] = None,
-                key: Optional[SeedOrKey] = None):
-    p = _check_py_seq(p)
-    jit_error_if(jnp.any(jnp.logical_and(p < 0, p > 1)), self._check_p, p)
-    if size is None:
-      size = jnp.shape(p)
-    key = self.split_key() if key is None else _formalize_key(key)
-    r = jr.bernoulli(key, p=p, shape=_size2shape(size))
-    return r
-
-  def lognormal(
-      self,
-      mean=None,
-      sigma=None,
-      size: Optional[Size] = None,
-      key: Optional[SeedOrKey] = None,
-      dtype: DTypeLike = None
-  ):
-    mean = _check_py_seq(mean)
-    sigma = _check_py_seq(sigma)
-    mean = bu.math.asarray(mean, dtype=dtype)
-    sigma = bu.math.asarray(sigma, dtype=dtype)
-    unit = mean.unit if isinstance(mean, bu.Quantity) else bu.Unit()
-    mean = mean.mantissa if isinstance(mean, bu.Quantity) else mean
-    sigma = sigma.in_unit(unit).mantissa if isinstance(sigma, bu.Quantity) else sigma
-
-    if size is None:
-      size = jnp.broadcast_shapes(
-        jnp.shape(mean),
-        jnp.shape(sigma)
-      )
-    key = self.split_key() if key is None else _formalize_key(key)
-    dtype = dtype or environ.dftype()
-    samples = jr.normal(key, shape=_size2shape(size), dtype=dtype)
-    samples = _loc_scale(mean, sigma, samples)
-    samples = jnp.exp(samples)
-    return samples if unit.is_unitless else bu.Quantity(samples, unit=unit)
-
-  def binomial(self,
-               n,
-               p,
-               size: Optional[Size] = None,
-               key: Optional[SeedOrKey] = None,
-               dtype: DTypeLike = None):
-    n = _check_py_seq(n)
-    p = _check_py_seq(p)
-    jit_error_if(jnp.any(jnp.logical_and(p < 0, p > 1)), self._check_p, p)
-    if size is None:
-      size = jnp.broadcast_shapes(jnp.shape(n), jnp.shape(p))
-    key = self.split_key() if key is None else _formalize_key(key)
-    r = _binomial(key, p, n, shape=_size2shape(size))
-    dtype = dtype or environ.ditype()
-    return jnp.asarray(r, dtype=dtype)
-
-  def chisquare(self,
-                df,
-                size: Optional[Size] = None,
-                key: Optional[SeedOrKey] = None,
-                dtype: DTypeLike = None):
-    df = _check_py_seq(df)
-    key = self.split_key() if key is None else _formalize_key(key)
-    dtype = dtype or environ.dftype()
-    if size is None:
-      if jnp.ndim(df) == 0:
-        dist = jr.normal(key, (df,), dtype=dtype) ** 2
-        dist = dist.sum()
-      else:
-        raise NotImplementedError('Do not support non-scale "df" when "size" is None')
-    else:
-      dist = jr.normal(key, (df,) + _size2shape(size), dtype=dtype) ** 2
-      dist = dist.sum(axis=0)
-    return dist
-
-  def dirichlet(self,
-                alpha,
-                size: Optional[Size] = None,
-                key: Optional[SeedOrKey] = None,
-                dtype: DTypeLike = None):
-    key = self.split_key() if key is None else _formalize_key(key)
-    alpha = _check_py_seq(alpha)
-    dtype = dtype or environ.dftype()
-    r = jr.dirichlet(key, alpha=alpha, shape=_size2shape(size), dtype=dtype)
-    return r
-
-  def geometric(self,
-                p,
-                size: Optional[Size] = None,
-                key: Optional[SeedOrKey] = None,
-                dtype: DTypeLike = None):
-    p = _check_py_seq(p)
-    if size is None:
-      size = jnp.shape(p)
-    key = self.split_key() if key is None else _formalize_key(key)
-    dtype = dtype or environ.dftype()
-    u = uniform_for_unit(key, size, dtype=dtype)
-    r = jnp.floor(jnp.log1p(-u) / jnp.log1p(-p))
-    return r
-
-  def _check_p2(self, p):
-    raise ValueError(f'We require `sum(pvals[:-1]) <= 1`. But we got {p}')
-
-  def multinomial(self,
-                  n,
-                  pvals,
-                  size: Optional[Size] = None,
-                  key: Optional[SeedOrKey] = None,
-                  dtype: DTypeLike = None):
-    key = self.split_key() if key is None else _formalize_key(key)
-    n = _check_py_seq(n)
-    pvals = _check_py_seq(pvals)
-    jit_error_if(jnp.sum(pvals[:-1]) > 1., self._check_p2, pvals)
-    if isinstance(n, jax.core.Tracer):
-      raise ValueError("The total count parameter `n` should not be a jax abstract array.")
-    size = _size2shape(size)
-    n_max = int(np.max(jax.device_get(n)))
-    batch_shape = lax.broadcast_shapes(jnp.shape(pvals)[:-1], jnp.shape(n))
-    r = _multinomial(key, pvals, n, n_max, batch_shape + size)
-    dtype = dtype or environ.ditype()
-    return jnp.asarray(r, dtype=dtype)
-
-  def multivariate_normal(
-      self,
-      mean,
-      cov,
-      size: Optional[Size] = None,
-      method: str = 'cholesky',
-      key: Optional[SeedOrKey] = None,
-      dtype: DTypeLike = None
-  ):
-    if method not in {'svd', 'eigh', 'cholesky'}:
-      raise ValueError("method must be one of {'svd', 'eigh', 'cholesky'}")
-    dtype = dtype or environ.dftype()
-    mean = bu.math.asarray(_check_py_seq(mean), dtype=dtype)
-    cov = bu.math.asarray(_check_py_seq(cov), dtype=dtype)
-    if isinstance(mean, bu.Quantity):
-      assert isinstance(cov, bu.Quantity)
-      assert mean.unit ** 2 == cov.unit
-    mean = mean.mantissa if isinstance(mean, bu.Quantity) else mean
-    cov = cov.mantissa if isinstance(cov, bu.Quantity) else cov
-    unit = mean.unit if isinstance(mean, bu.Quantity) else bu.Unit()
-
-    key = self.split_key() if key is None else _formalize_key(key)
-    if not jnp.ndim(mean) >= 1:
-      raise ValueError(f"multivariate_normal requires mean.ndim >= 1, got mean.ndim == {jnp.ndim(mean)}")
-    if not jnp.ndim(cov) >= 2:
-      raise ValueError(f"multivariate_normal requires cov.ndim >= 2, got cov.ndim == {jnp.ndim(cov)}")
-    n = mean.shape[-1]
-    if jnp.shape(cov)[-2:] != (n, n):
-      raise ValueError(f"multivariate_normal requires cov.shape == (..., n, n) for n={n}, "
-                       f"but got cov.shape == {jnp.shape(cov)}.")
-    if size is None:
-      size = lax.broadcast_shapes(mean.shape[:-1], cov.shape[:-2])
-    else:
-      size = _size2shape(size)
-      _check_shape("normal", size, mean.shape[:-1], cov.shape[:-2])
-
-    if method == 'svd':
-      (u, s, _) = jnp.linalg.svd(cov)
-      factor = u * jnp.sqrt(s[..., None, :])
-    elif method == 'eigh':
-      (w, v) = jnp.linalg.eigh(cov)
-      factor = v * jnp.sqrt(w[..., None, :])
-    else:  # 'cholesky'
-      factor = jnp.linalg.cholesky(cov)
-    normal_samples = jr.normal(key, size + mean.shape[-1:], dtype=dtype)
-    r = mean + jnp.einsum('...ij,...j->...i', factor, normal_samples)
-    return r if unit.is_unitless else bu.Quantity(r, unit=unit)
-
-  def rayleigh(self,
-               scale=1.0,
-               size: Optional[Size] = None,
-               key: Optional[SeedOrKey] = None,
-               dtype: DTypeLike = None):
-    scale = _check_py_seq(scale)
-    if size is None:
-      size = jnp.shape(scale)
-    key = self.split_key() if key is None else _formalize_key(key)
-    dtype = dtype or environ.dftype()
-    x = jnp.sqrt(-2. * jnp.log(uniform_for_unit(key, shape=_size2shape(size), minval=0, maxval=1, dtype=dtype)))
-    r = x * scale
-    return r
-
-  def triangular(self,
-                 size: Optional[Size] = None,
-                 key: Optional[SeedOrKey] = None):
-    key = self.split_key() if key is None else _formalize_key(key)
-    bernoulli_samples = jr.bernoulli(key, p=0.5, shape=_size2shape(size))
-    r = 2 * bernoulli_samples - 1
-    return r
-
-  def vonmises(self,
-               mu,
-               kappa,
-               size: Optional[Size] = None,
-               key: Optional[SeedOrKey] = None,
-               dtype: DTypeLike = None):
-    key = self.split_key() if key is None else _formalize_key(key)
-    dtype = dtype or environ.dftype()
-    mu = jnp.asarray(_check_py_seq(mu), dtype=dtype)
-    kappa = jnp.asarray(_check_py_seq(kappa), dtype=dtype)
-    if size is None:
-      size = lax.broadcast_shapes(jnp.shape(mu), jnp.shape(kappa))
-    size = _size2shape(size)
-    samples = _von_mises_centered(key, kappa, size, dtype=dtype)
-    samples = samples + mu
-    samples = (samples + jnp.pi) % (2.0 * jnp.pi) - jnp.pi
-    return samples
-
-  def weibull(self,
-              a,
-              size: Optional[Size] = None,
-              key: Optional[SeedOrKey] = None,
-              dtype: DTypeLike = None):
-    key = self.split_key() if key is None else _formalize_key(key)
-    a = _check_py_seq(a)
-    if size is None:
-      size = jnp.shape(a)
-    else:
-      if jnp.size(a) > 1:
-        raise ValueError(f'"a" should be a scalar when "size" is provided. But we got {a}')
-    size = _size2shape(size)
-    dtype = dtype or environ.dftype()
-    random_uniform = uniform_for_unit(key=key, shape=size, minval=0, maxval=1, dtype=dtype)
-    r = jnp.power(-jnp.log1p(-random_uniform), 1.0 / a)
-    return r
-
-  def weibull_min(self,
-                  a,
-                  scale=None,
-                  size: Optional[Size] = None,
-                  key: Optional[SeedOrKey] = None,
-                  dtype: DTypeLike = None):
-    key = self.split_key() if key is None else _formalize_key(key)
-    a = _check_py_seq(a)
-    scale = _check_py_seq(scale)
-    if size is None:
-      size = jnp.broadcast_shapes(jnp.shape(a), jnp.shape(scale))
-    else:
-      if jnp.size(a) > 1:
-        raise ValueError(f'"a" should be a scalar when "size" is provided. But we got {a}')
-    size = _size2shape(size)
-    dtype = dtype or environ.dftype()
-    random_uniform = uniform_for_unit(key=key, shape=size, minval=0, maxval=1, dtype=dtype)
-    r = jnp.power(-jnp.log1p(-random_uniform), 1.0 / a)
-    if scale is not None:
-      r /= scale
-    return r
-
-  def maxwell(self,
-              size: Optional[Size] = None,
-              key: Optional[SeedOrKey] = None,
-              dtype: DTypeLike = None):
-    key = self.split_key() if key is None else _formalize_key(key)
-    shape = _size2shape(size) + (3,)
-    dtype = dtype or environ.dftype()
-    norm_rvs = jr.normal(key=key, shape=shape, dtype=dtype)
-    r = jnp.linalg.norm(norm_rvs, axis=-1)
-    return r
-
-  def negative_binomial(self,
-                        n,
-                        p,
-                        size: Optional[Size] = None,
-                        key: Optional[SeedOrKey] = None,
-                        dtype: DTypeLike = None):
-    n = _check_py_seq(n)
-    p = _check_py_seq(p)
-    if size is None:
-      size = lax.broadcast_shapes(jnp.shape(n), jnp.shape(p))
-    size = _size2shape(size)
-    logits = jnp.log(p) - jnp.log1p(-p)
-    if key is None:
-      keys = self.split_keys(2)
-    else:
-      keys = jr.split(_formalize_key(key), 2)
-    rate = self.gamma(shape=n, scale=jnp.exp(-logits), size=size, key=keys[0], dtype=environ.dftype())
-    r = self.poisson(lam=rate, key=keys[1], dtype=dtype or environ.ditype())
-    return r
-
-  def wald(self,
-           mean,
-           scale,
-           size: Optional[Size] = None,
-           key: Optional[SeedOrKey] = None,
-           dtype: DTypeLike = None):
-    dtype = dtype or environ.dftype()
-    key = self.split_key() if key is None else _formalize_key(key)
-    mean = jnp.asarray(_check_py_seq(mean), dtype=dtype)
-    scale = jnp.asarray(_check_py_seq(scale), dtype=dtype)
-    if size is None:
-      size = lax.broadcast_shapes(jnp.shape(mean), jnp.shape(scale))
-    size = _size2shape(size)
-    sampled_chi2 = jnp.square(self.randn(*size))
-    sampled_uniform = self.uniform(size=size, key=key, dtype=dtype)
-    # Wikipedia defines an intermediate x with the formula
-    #   x = loc + loc ** 2 * y / (2 * conc) - loc / (2 * conc) * sqrt(4 * loc * conc * y + loc ** 2 * y ** 2)
-    # where y ~ N(0, 1)**2 (sampled_chi2 above) and conc is the concentration.
-    # Let us write
-    #   w = loc * y / (2 * conc)
-    # Then we can extract the common factor in the last two terms to obtain
-    #   x = loc + loc * w * (1 - sqrt(2 / w + 1))
-    # Now we see that the Wikipedia formula suffers from catastrphic
-    # cancellation for large w (e.g., if conc << loc).
-    #
-    # Fortunately, we can fix this by multiplying both sides
-    # by 1 + sqrt(2 / w + 1).  We get
-    #   x * (1 + sqrt(2 / w + 1)) =
-    #     = loc * (1 + sqrt(2 / w + 1)) + loc * w * (1 - (2 / w + 1))
-    #     = loc * (sqrt(2 / w + 1) - 1)
-    # The term sqrt(2 / w + 1) + 1 no longer presents numerical
-    # difficulties for large w, and sqrt(2 / w + 1) - 1 is just
-    # sqrt1pm1(2 / w), which we know how to compute accurately.
-    # This just leaves the matter of small w, where 2 / w may
-    # overflow.  In the limit a w -> 0, x -> loc, so we just mask
-    # that case.
-    sqrt1pm1_arg = 4 * scale / (mean * sampled_chi2)  # 2 / w above
-    safe_sqrt1pm1_arg = jnp.where(sqrt1pm1_arg < np.inf, sqrt1pm1_arg, 1.0)
-    denominator = 1.0 + jnp.sqrt(safe_sqrt1pm1_arg + 1.0)
-    ratio = jnp.expm1(0.5 * jnp.log1p(safe_sqrt1pm1_arg)) / denominator
-    sampled = mean * jnp.where(sqrt1pm1_arg < np.inf, ratio, 1.0)  # x above
-    res = jnp.where(sampled_uniform <= mean / (mean + sampled),
-                    sampled,
-                    jnp.square(mean) / sampled)
-    return res
-
-  def t(self,
-        df,
-        size: Optional[Size] = None,
-        key: Optional[SeedOrKey] = None,
-        dtype: DTypeLike = None):
-    dtype = dtype or environ.dftype()
-    df = jnp.asarray(_check_py_seq(df), dtype=dtype)
-    if size is None:
-      size = np.shape(df)
-    else:
-      size = _size2shape(size)
-      _check_shape("t", size, np.shape(df))
-    if key is None:
-      keys = self.split_keys(2)
-    else:
-      keys = jr.split(_formalize_key(key), 2)
-    n = jr.normal(keys[0], size, dtype=dtype)
-    two = _const(n, 2)
-    half_df = lax.div(df, two)
-    g = jr.gamma(keys[1], half_df, size, dtype=dtype)
-    r = n * jnp.sqrt(half_df / g)
-    return r
-
-  def orthogonal(self,
-                 n: int,
-                 size: Optional[Size] = None,
-                 key: Optional[SeedOrKey] = None,
-                 dtype: DTypeLike = None):
-    dtype = dtype or environ.dftype()
-    key = self.split_key() if key is None else _formalize_key(key)
-    size = _size2shape(size)
-    _check_shape("orthogonal", size)
-    n = core.concrete_or_error(index, n, "The error occurred in jax.random.orthogonal()")
-    z = jr.normal(key, size + (n, n), dtype=dtype)
-    q, r = jnp.linalg.qr(z)
-    d = jnp.diagonal(r, 0, -2, -1)
-    r = q * jnp.expand_dims(d / abs(d), -2)
-    return r
-
-  def noncentral_chisquare(self,
-                           df,
-                           nonc,
-                           size: Optional[Size] = None,
-                           key: Optional[SeedOrKey] = None,
-                           dtype: DTypeLike = None):
-    dtype = dtype or environ.dftype()
-    df = jnp.asarray(_check_py_seq(df), dtype=dtype)
-    nonc = jnp.asarray(_check_py_seq(nonc), dtype=dtype)
-    if size is None:
-      size = lax.broadcast_shapes(jnp.shape(df), jnp.shape(nonc))
-    size = _size2shape(size)
-    if key is None:
-      keys = self.split_keys(3)
-    else:
-      keys = jr.split(_formalize_key(key), 3)
-    i = jr.poisson(keys[0], 0.5 * nonc, shape=size, dtype=environ.ditype())
-    n = jr.normal(keys[1], shape=size, dtype=dtype) + jnp.sqrt(nonc)
-    cond = jnp.greater(df, 1.0)
-    df2 = jnp.where(cond, df - 1.0, df + 2.0 * i)
-    chi2 = 2.0 * jr.gamma(keys[2], 0.5 * df2, shape=size, dtype=dtype)
-    r = jnp.where(cond, chi2 + n * n, chi2)
-    return r
-
-  def loggamma(self,
-               a,
-               size: Optional[Size] = None,
-               key: Optional[SeedOrKey] = None,
-               dtype: DTypeLike = None):
-    dtype = dtype or environ.dftype()
-    key = self.split_key() if key is None else _formalize_key(key)
-    a = _check_py_seq(a)
-    if size is None:
-      size = jnp.shape(a)
-    r = jr.loggamma(key, a, shape=_size2shape(size), dtype=dtype)
-    return r
-
-  def categorical(self,
-                  logits,
-                  axis: int = -1,
-                  size: Optional[Size] = None,
-                  key: Optional[SeedOrKey] = None):
-    key = self.split_key() if key is None else _formalize_key(key)
-    logits = _check_py_seq(logits)
-    if size is None:
-      size = list(jnp.shape(logits))
-      size.pop(axis)
-    r = jr.categorical(key, logits, axis=axis, shape=_size2shape(size))
-    return r
-
-  def zipf(self,
-           a,
-           size: Optional[Size] = None,
-           key: Optional[SeedOrKey] = None,
-           dtype: DTypeLike = None):
-    a = _check_py_seq(a)
-    if size is None:
-      size = jnp.shape(a)
-    dtype = dtype or environ.ditype()
-    r = jax.pure_callback(lambda x: np.random.zipf(x, size).astype(dtype),
-                          jax.ShapeDtypeStruct(size, dtype),
-                          a)
-    return r
-
-  def power(self,
-            a,
-            size: Optional[Size] = None,
-            key: Optional[SeedOrKey] = None,
-            dtype: DTypeLike = None):
-    a = _check_py_seq(a)
-    if size is None:
-      size = jnp.shape(a)
-    size = _size2shape(size)
-    dtype = dtype or environ.dftype()
-    r = jax.pure_callback(lambda a: np.random.power(a=a, size=size).astype(dtype),
-                          jax.ShapeDtypeStruct(size, dtype),
-                          a)
-    return r
-
-  def f(self,
-        dfnum,
-        dfden,
-        size: Optional[Size] = None,
-        key: Optional[SeedOrKey] = None,
-        dtype: DTypeLike = None):
-    dfnum = _check_py_seq(dfnum)
-    dfden = _check_py_seq(dfden)
-    if size is None:
-      size = jnp.broadcast_shapes(jnp.shape(dfnum), jnp.shape(dfden))
-    size = _size2shape(size)
-    d = {'dfnum': dfnum, 'dfden': dfden}
-    dtype = dtype or environ.dftype()
-    r = jax.pure_callback(lambda dfnum_, dfden_: np.random.f(dfnum=dfnum_,
-                                                             dfden=dfden_,
-                                                             size=size).astype(dtype),
-                          jax.ShapeDtypeStruct(size, dtype),
-                          dfnum, dfden)
-    return r
-
-  def hypergeometric(
-      self,
-      ngood,
-      nbad,
-      nsample,
-      size: Optional[Size] = None,
-      key: Optional[SeedOrKey] = None,
-      dtype: DTypeLike = None
-  ):
-    ngood = _check_py_seq(ngood)
-    nbad = _check_py_seq(nbad)
-    nsample = _check_py_seq(nsample)
-
-    if size is None:
-      size = lax.broadcast_shapes(jnp.shape(ngood),
-                                  jnp.shape(nbad),
-                                  jnp.shape(nsample))
-    size = _size2shape(size)
-    dtype = dtype or environ.ditype()
-    d = {'ngood': ngood, 'nbad': nbad, 'nsample': nsample}
-    r = jax.pure_callback(lambda d: np.random.hypergeometric(ngood=d['ngood'],
-                                                             nbad=d['nbad'],
-                                                             nsample=d['nsample'],
-                                                             size=size).astype(dtype),
-                          jax.ShapeDtypeStruct(size, dtype),
-                          d)
-    return r
-
-  def logseries(self,
-                p,
-                size: Optional[Size] = None,
-                key: Optional[SeedOrKey] = None,
-                dtype: DTypeLike = None):
-    p = _check_py_seq(p)
-    if size is None:
-      size = jnp.shape(p)
-    size = _size2shape(size)
-    dtype = dtype or environ.ditype()
-    r = jax.pure_callback(lambda p: np.random.logseries(p=p, size=size).astype(dtype),
-                          jax.ShapeDtypeStruct(size, dtype),
-                          p)
-    return r
-
-  def noncentral_f(self,
-                   dfnum,
-                   dfden,
-                   nonc,
-                   size: Optional[Size] = None,
-                   key: Optional[SeedOrKey] = None,
-                   dtype: DTypeLike = None):
-    dfnum = _check_py_seq(dfnum)
-    dfden = _check_py_seq(dfden)
-    nonc = _check_py_seq(nonc)
-    if size is None:
-      size = lax.broadcast_shapes(jnp.shape(dfnum),
-                                  jnp.shape(dfden),
-                                  jnp.shape(nonc))
-    size = _size2shape(size)
-    d = {'dfnum': dfnum, 'dfden': dfden, 'nonc': nonc}
-    dtype = dtype or environ.dftype()
-    r = jax.pure_callback(lambda x: np.random.noncentral_f(dfnum=x['dfnum'],
-                                                           dfden=x['dfden'],
-                                                           nonc=x['nonc'],
-                                                           size=size).astype(dtype),
-                          jax.ShapeDtypeStruct(size, dtype),
-                          d)
-    return r
-
-  # PyTorch compatibility #
-  # --------------------- #
-
-  def rand_like(self, input, *, dtype=None, key: Optional[SeedOrKey] = None):
-    """Returns a tensor with the same size as input that is filled with random
-    numbers from a uniform distribution on the interval ``[0, 1)``.
-
-    Args:
-      input:  the ``size`` of input will determine size of the output tensor.
-      dtype:  the desired data type of returned Tensor. Default: if ``None``, defaults to the dtype of input.
-      key: the seed or key for the random.
-
-    Returns:
-      The random data.
-    """
-    return self.random(jnp.shape(input), key=key).astype(dtype)
-
-  def randn_like(self, input, *, dtype=None, key: Optional[SeedOrKey] = None):
-    """Returns a tensor with the same size as ``input`` that is filled with
-    random numbers from a normal distribution with mean 0 and variance 1.
-
-    Args:
-      input:  the ``size`` of input will determine size of the output tensor.
-      dtype:  the desired data type of returned Tensor. Default: if ``None``, defaults to the dtype of input.
-      key: the seed or key for the random.
-
-    Returns:
-      The random data.
-    """
-    return self.randn(*jnp.shape(input), key=key).astype(dtype)
-
-  def randint_like(self, input, low=0, high=None, *, dtype=None, key: Optional[SeedOrKey] = None):
-    if high is None:
-      high = max(input)
-    return self.randint(low, high=high, size=jnp.shape(input), dtype=dtype, key=key)
-
-
-# default random generator
-DEFAULT = RandomState(np.random.randint(0, 10000, size=2, dtype=np.uint32))
-
-
-def split_key():
-  """Create a new seed from the current seed.
-
-  This function is useful for the consistency with JAX's random paradigm."""
-  return DEFAULT.split_key()
-
-
-def split_keys(n):
-  """Create multiple seeds from the current seed. This is used
-  internally by `pmap` and `vmap` to ensure that random numbers
-  are different in parallel threads.
-
-  Parameters
-  ----------
-  n : int
-    The number of seeds to generate.
-  """
-  return DEFAULT.split_keys(n)
-
-
-def clone_rng(seed_or_key=None, clone: bool = True) -> RandomState:
-  """Clone the random state according to the given setting.
-
-  Args:
-    seed_or_key: The seed (an integer) or the random key.
-    clone: Bool. Whether clone the default random state.
-
-  Returns:
-    The random state.
-  """
-  if seed_or_key is None:
-    return DEFAULT.clone() if clone else DEFAULT
-  else:
-    return RandomState(seed_or_key)
-
-
-def default_rng(seed_or_key=None, clone: bool = True) -> RandomState:
-  """
-  Get the default random state.
-
-  Args:
-    seed_or_key: The seed (an integer) or the jax random key.
-    clone: Bool. Whether clone the default random state.
-
-  Returns:
-    The random state.
-  """
-  if seed_or_key is None:
-    return DEFAULT.clone() if clone else DEFAULT
-  else:
-    return RandomState(seed_or_key)
-
-
-def seed(seed_or_key: int = None):
-  """Sets a new random seed.
-
-  Parameters
-  ----------
-  seed_or_key: int, optional
-    The random seed (an integer) or jax random key.
-  """
-  with jax.ensure_compile_time_eval():
-    if seed_or_key is None:
-      seed_or_key = np.random.randint(0, 100000)
-
-  # numpy random seed
-  if np.size(seed_or_key) == 1:  # seed
-    np.random.seed(seed_or_key)
-  elif np.size(seed_or_key) == 2:  # jax random key
-    np.random.seed(seed_or_key[0])
-  else:
-    raise ValueError(f"seed_or_key should be an integer or a tuple of two integers.")
-
-  # jax random seed
-  DEFAULT.seed(seed_or_key)
-
-
-@contextmanager
-def seed_context(seed_or_key: SeedOrKey):
-  """
-  A context manager that sets the random seed for the duration of the block.
-
-  Examples:
-
-  >>> import brainstate as bst
-  >>> print(bst.random.rand(2))
-  [0.57721865 0.9820676 ]
-  >>> print(bst.random.rand(2))
-  [0.8511752  0.95312667]
-  >>> with bst.random.seed_context(42):
-  ...     print(bst.random.rand(2))
-  [0.95598125 0.4032725 ]
-  >>> with bst.random.seed_context(42):
-  ...     print(bst.random.rand(2))
-  [0.95598125 0.4032725 ]
-
-  .. note::
-
-     The context manager does not only set the seed for the AX random state, but also for the numpy random state.
-
-  Args:
-    seed_or_key: The seed (an integer) or jax random key.
-
-  """
-  old_jrand_key = DEFAULT.value
-  old_np_state = np.random.get_state()
-  try:
-    if np.size(seed_or_key) == 1:  # seed
-      np.random.seed(seed_or_key)
-    elif np.size(seed_or_key) == 2:  # jax random key
-      np.random.seed(seed_or_key[0])
-    else:
-      raise ValueError(f"seed_or_key should be an integer or a tuple of two integers.")
-    DEFAULT.seed(seed_or_key)
-    yield
-  finally:
-    np.random.set_state(old_np_state)
-    DEFAULT.seed(old_jrand_key)
-
-
 def rand(*dn, key: Optional[SeedOrKey] = None, dtype: DTypeLike = None):
   r"""
   Random values in a given shape.
@@ -1260,6 +63,9 @@ def rand(*dn, key: Optional[SeedOrKey] = None, dtype: DTypeLike = None):
   dtype : dtype, optional
       Desired dtype of the result. Byteorder must be native.
       The default value is float.
+  key : PRNGKey, optional
+      The key for the random number generator. If not given, the
+      default random number generator is used.
 
   Returns
   -------
@@ -1272,7 +78,8 @@ def rand(*dn, key: Optional[SeedOrKey] = None, dtype: DTypeLike = None):
 
   Examples
   --------
-  >>> brainstate.random.rand(3,2)
+  >>> import brainstate as bst
+  >>> bst.random.rand(3,2)
   array([[ 0.14022471,  0.96360618],  #random
          [ 0.37601032,  0.25528411],  #random
          [ 0.49313049,  0.94909878]]) #random
@@ -4790,402 +3597,6 @@ def randint_like(input, low=0, high=None, *, dtype=None, key: Optional[SeedOrKey
 # ---------------------------------------------------------------------------------------------------------------
 
 
-def _formalize_key(key):
-  if isinstance(key, int):
-    return jr.PRNGKey(key)
-  elif isinstance(key, (jax.Array, np.ndarray)):
-    if key.dtype != jnp.uint32:
-      raise TypeError('key must be a int or an array with two uint32.')
-    if key.size != 2:
-      raise TypeError('key must be a int or an array with two uint32.')
-    return jnp.asarray(key, dtype=jnp.uint32)
-  else:
-    raise TypeError('key must be a int or an array with two uint32.')
-
-
-def _size2shape(size):
-  if size is None:
-    return ()
-  elif isinstance(size, (tuple, list)):
-    return tuple(size)
-  else:
-    return (size,)
-
-
-def _check_shape(name, shape, *param_shapes):
-  if param_shapes:
-    shape_ = lax.broadcast_shapes(shape, *param_shapes)
-    if shape != shape_:
-      msg = ("{} parameter shapes must be broadcast-compatible with shape "
-             "argument, and the result of broadcasting the shapes must equal "
-             "the shape argument, but got result {} for shape argument {}.")
-      raise ValueError(msg.format(name, shape_, shape))
-
-
-def _is_python_scalar(x):
-  if hasattr(x, 'aval'):
-    return x.aval.weak_type
-  elif np.ndim(x) == 0:
-    return True
-  elif isinstance(x, (bool, int, float, complex)):
-    return True
-  else:
-    return False
-
-
-python_scalar_dtypes = {
-  bool: np.dtype('bool'),
-  int: np.dtype('int64'),
-  float: np.dtype('float64'),
-  complex: np.dtype('complex128'),
-}
-
-
-def _dtype(x, *, canonicalize: bool = False):
-  """Return the dtype object for a value or type, optionally canonicalized based on X64 mode."""
-  if x is None:
-    raise ValueError(f"Invalid argument to dtype: {x}.")
-  elif isinstance(x, type) and x in python_scalar_dtypes:
-    dt = python_scalar_dtypes[x]
-  elif type(x) in python_scalar_dtypes:
-    dt = python_scalar_dtypes[type(x)]
-  elif hasattr(x, 'dtype'):
-    dt = x.dtype
-  else:
-    dt = np.result_type(x)
-  return dtypes.canonicalize_dtype(dt) if canonicalize else dt
-
-
-def _const(example, val):
-  if _is_python_scalar(example):
-    dtype = dtypes.canonicalize_dtype(type(example))
-    val = dtypes.scalar_type_of(example)(val)
-    return val if dtype == _dtype(val, canonicalize=True) else np.array(val, dtype)
-  else:
-    dtype = dtypes.canonicalize_dtype(example.dtype)
-  return np.array(val, dtype)
-
-
-_tr_params = namedtuple(
-  "tr_params", ["c", "b", "a", "alpha", "u_r", "v_r", "m", "log_p", "log1_p", "log_h"]
-)
-
-
-def _get_tr_params(n, p):
-  # See Table 1. Additionally, we pre-compute log(p), log1(-p) and the
-  # constant terms, that depend only on (n, p, m) in log(f(k)) (bottom of page 5).
-  mu = n * p
-  spq = jnp.sqrt(mu * (1 - p))
-  c = mu + 0.5
-  b = 1.15 + 2.53 * spq
-  a = -0.0873 + 0.0248 * b + 0.01 * p
-  alpha = (2.83 + 5.1 / b) * spq
-  u_r = 0.43
-  v_r = 0.92 - 4.2 / b
-  m = jnp.floor((n + 1) * p).astype(n.dtype)
-  log_p = jnp.log(p)
-  log1_p = jnp.log1p(-p)
-  log_h = ((m + 0.5) * (jnp.log((m + 1.0) / (n - m + 1.0)) + log1_p - log_p) +
-           _stirling_approx_tail(m) + _stirling_approx_tail(n - m))
-  return _tr_params(c, b, a, alpha, u_r, v_r, m, log_p, log1_p, log_h)
-
-
-def _stirling_approx_tail(k):
-  precomputed = jnp.array([0.08106146679532726,
-                           0.04134069595540929,
-                           0.02767792568499834,
-                           0.02079067210376509,
-                           0.01664469118982119,
-                           0.01387612882307075,
-                           0.01189670994589177,
-                           0.01041126526197209,
-                           0.009255462182712733,
-                           0.008330563433362871],
-                          dtype=environ.dftype())
-  kp1 = k + 1
-  kp1sq = (k + 1) ** 2
-  return jnp.where(k < 10,
-                   precomputed[k],
-                   (1.0 / 12 - (1.0 / 360 - (1.0 / 1260) / kp1sq) / kp1sq) / kp1)
-
-
-def _binomial_btrs(key, p, n):
-  """
-  Based on the transformed rejection sampling algorithm (BTRS) from the
-  following reference:
-
-  Hormann, "The Generation of Binonmial Random Variates"
-  (https://core.ac.uk/download/pdf/11007254.pdf)
-  """
-
-  def _btrs_body_fn(val):
-    _, key, _, _ = val
-    key, key_u, key_v = jr.split(key, 3)
-    u = jr.uniform(key_u)
-    v = jr.uniform(key_v)
-    u = u - 0.5
-    k = jnp.floor(
-      (2 * tr_params.a / (0.5 - jnp.abs(u)) + tr_params.b) * u + tr_params.c
-    ).astype(n.dtype)
-    return k, key, u, v
-
-  def _btrs_cond_fn(val):
-    def accept_fn(k, u, v):
-      # See acceptance condition in Step 3. (Page 3) of TRS algorithm
-      # v <= f(k) * g_grad(u) / alpha
-
-      m = tr_params.m
-      log_p = tr_params.log_p
-      log1_p = tr_params.log1_p
-      # See: formula for log(f(k)) at bottom of Page 5.
-      log_f = (
-          (n + 1.0) * jnp.log((n - m + 1.0) / (n - k + 1.0))
-          + (k + 0.5) * (jnp.log((n - k + 1.0) / (k + 1.0)) + log_p - log1_p)
-          + (_stirling_approx_tail(k) - _stirling_approx_tail(n - k))
-          + tr_params.log_h
-      )
-      g = (tr_params.a / (0.5 - jnp.abs(u)) ** 2) + tr_params.b
-      return jnp.log((v * tr_params.alpha) / g) <= log_f
-
-    k, key, u, v = val
-    early_accept = (jnp.abs(u) <= tr_params.u_r) & (v <= tr_params.v_r)
-    early_reject = (k < 0) | (k > n)
-    return lax.cond(
-      early_accept | early_reject,
-      (),
-      lambda _: ~early_accept,
-      (k, u, v),
-      lambda x: ~accept_fn(*x),
-    )
-
-  tr_params = _get_tr_params(n, p)
-  ret = lax.while_loop(
-    _btrs_cond_fn, _btrs_body_fn, (-1, key, 1.0, 1.0)
-  )  # use k=-1 initially so that cond_fn returns True
-  return ret[0]
-
-
-def _binomial_inversion(key, p, n):
-  def _binom_inv_body_fn(val):
-    i, key, geom_acc = val
-    key, key_u = jr.split(key)
-    u = jr.uniform(key_u)
-    geom = jnp.floor(jnp.log1p(-u) / log1_p) + 1
-    geom_acc = geom_acc + geom
-    return i + 1, key, geom_acc
-
-  def _binom_inv_cond_fn(val):
-    i, _, geom_acc = val
-    return geom_acc <= n
-
-  log1_p = jnp.log1p(-p)
-  ret = lax.while_loop(_binom_inv_cond_fn, _binom_inv_body_fn, (-1, key, 0.0))
-  return ret[0]
-
-
-def _binomial_dispatch(key, p, n):
-  def dispatch(key, p, n):
-    is_le_mid = p <= 0.5
-    pq = jnp.where(is_le_mid, p, 1 - p)
-    mu = n * pq
-    k = lax.cond(
-      mu < 10,
-      (key, pq, n),
-      lambda x: _binomial_inversion(*x),
-      (key, pq, n),
-      lambda x: _binomial_btrs(*x),
-    )
-    return jnp.where(is_le_mid, k, n - k)
-
-  # Return 0 for nan `p` or negative `n`, since nan values are not allowed for integer types
-  cond0 = jnp.isfinite(p) & (n > 0) & (p > 0)
-  return lax.cond(
-    cond0 & (p < 1),
-    (key, p, n),
-    lambda x: dispatch(*x),
-    (),
-    lambda _: jnp.where(cond0, n, 0),
-  )
-
-
-@partial(jit, static_argnums=(3,))
-def _binomial(key, p, n, shape):
-  shape = shape or lax.broadcast_shapes(jnp.shape(p), jnp.shape(n))
-  # reshape to map over axis 0
-  p = jnp.reshape(jnp.broadcast_to(p, shape), -1)
-  n = jnp.reshape(jnp.broadcast_to(n, shape), -1)
-  key = jr.split(key, jnp.size(p))
-  if jax.default_backend() == "cpu":
-    ret = lax.map(lambda x: _binomial_dispatch(*x), (key, p, n))
-  else:
-    ret = vmap(lambda *x: _binomial_dispatch(*x))(key, p, n)
-  return jnp.reshape(ret, shape)
-
-
-@partial(jit, static_argnums=(2,))
-def _categorical(key, p, shape):
-  # this implementation is fast when event shape is small, and slow otherwise
-  # Ref: https://stackoverflow.com/a/34190035
-  shape = shape or p.shape[:-1]
-  s = jnp.cumsum(p, axis=-1)
-  r = jr.uniform(key, shape=shape + (1,))
-  return jnp.sum(s < r, axis=-1)
-
-
-def _scatter_add_one(operand, indices, updates):
-  return lax.scatter_add(
-    operand,
-    indices,
-    updates,
-    lax.ScatterDimensionNumbers(
-      update_window_dims=(),
-      inserted_window_dims=(0,),
-      scatter_dims_to_operand_dims=(0,),
-    ),
-  )
-
-
-def _reshape(x, shape):
-  if isinstance(x, (int, float, np.ndarray, np.generic)):
-    return np.reshape(x, shape)
-  else:
-    return jnp.reshape(x, shape)
-
-
-def _promote_shapes(*args, shape=()):
-  # adapted from lax.lax_numpy
-  if len(args) < 2 and not shape:
-    return args
-  else:
-    shapes = [jnp.shape(arg) for arg in args]
-    num_dims = len(lax.broadcast_shapes(shape, *shapes))
-    return [
-      _reshape(arg, (1,) * (num_dims - len(s)) + s) if len(s) < num_dims else arg
-      for arg, s in zip(args, shapes)
-    ]
-
-
-@partial(jit, static_argnums=(3, 4))
-def _multinomial(key, p, n, n_max, shape=()):
-  if jnp.shape(n) != jnp.shape(p)[:-1]:
-    broadcast_shape = lax.broadcast_shapes(jnp.shape(n), jnp.shape(p)[:-1])
-    n = jnp.broadcast_to(n, broadcast_shape)
-    p = jnp.broadcast_to(p, broadcast_shape + jnp.shape(p)[-1:])
-  shape = shape or p.shape[:-1]
-  if n_max == 0:
-    return jnp.zeros(shape + p.shape[-1:], dtype=jnp.result_type(int))
-  # get indices from categorical distribution then gather the result
-  indices = _categorical(key, p, (n_max,) + shape)
-  # mask out values when counts is heterogeneous
-  if jnp.ndim(n) > 0:
-    mask = _promote_shapes(jnp.arange(n_max) < jnp.expand_dims(n, -1), shape=shape + (n_max,))[0]
-    mask = jnp.moveaxis(mask, -1, 0).astype(indices.dtype)
-    excess = jnp.concatenate([jnp.expand_dims(n_max - n, -1),
-                              jnp.zeros(jnp.shape(n) + (p.shape[-1] - 1,))],
-                             -1)
-  else:
-    mask = 1
-    excess = 0
-  # NB: we transpose to move batch shape to the front
-  indices_2D = (jnp.reshape(indices * mask, (n_max, -1))).T
-  samples_2D = vmap(_scatter_add_one)(jnp.zeros((indices_2D.shape[0], p.shape[-1]), dtype=indices.dtype),
-                                      jnp.expand_dims(indices_2D, axis=-1),
-                                      jnp.ones(indices_2D.shape, dtype=indices.dtype))
-  return jnp.reshape(samples_2D, shape + p.shape[-1:]) - excess
-
-
-@partial(jit, static_argnums=(2, 3), static_argnames=['shape', 'dtype'])
-def _von_mises_centered(key, concentration, shape, dtype=None):
-  """Compute centered von Mises samples using rejection sampling from [1]_ with wrapped Cauchy proposal.
-
-  Returns
-  -------
-  out: array_like
-     centered samples from von Mises
-
-  References
-  ----------
-  .. [1] Luc Devroye "Non-Uniform Random Variate Generation", Springer-Verlag, 1986;
-         Chapter 9, p. 473-476. http://www.nrbook.com/devroye/Devroye_files/chapter_nine.pdf
-
-  """
-  shape = shape or jnp.shape(concentration)
-  dtype = dtype or environ.dftype()
-  concentration = lax.convert_element_type(concentration, dtype)
-  concentration = jnp.broadcast_to(concentration, shape)
-
-  if dtype == jnp.float16:
-    s_cutoff = 1.8e-1
-  elif dtype == jnp.float32:
-    s_cutoff = 2e-2
-  elif dtype == jnp.float64:
-    s_cutoff = 1.2e-4
-  else:
-    raise ValueError(f"Unsupported dtype: {dtype}")
-
-  r = 1.0 + jnp.sqrt(1.0 + 4.0 * concentration ** 2)
-  rho = (r - jnp.sqrt(2.0 * r)) / (2.0 * concentration)
-  s_exact = (1.0 + rho ** 2) / (2.0 * rho)
-
-  s_approximate = 1.0 / concentration
-
-  s = jnp.where(concentration > s_cutoff, s_exact, s_approximate)
-
-  def cond_fn(*args):
-    """check if all are done or reached max number of iterations"""
-    i, _, done, _, _ = args[0]
-    return jnp.bitwise_and(i < 100, jnp.logical_not(jnp.all(done)))
-
-  def body_fn(*args):
-    i, key, done, _, w = args[0]
-    uni_ukey, uni_vkey, key = jr.split(key, 3)
-    u = jr.uniform(
-      key=uni_ukey,
-      shape=shape,
-      dtype=concentration.dtype,
-      minval=-1.0,
-      maxval=1.0,
-    )
-    z = jnp.cos(jnp.pi * u)
-    w = jnp.where(done, w, (1.0 + s * z) / (s + z))  # Update where not done
-    y = concentration * (s - w)
-    v = jr.uniform(key=uni_vkey, shape=shape, dtype=concentration.dtype)
-    accept = (y * (2.0 - y) >= v) | (jnp.log(y / v) + 1.0 >= y)
-    return i + 1, key, accept | done, u, w
-
-  init_done = jnp.zeros(shape, dtype=bool)
-  init_u = jnp.zeros(shape)
-  init_w = jnp.zeros(shape)
-
-  _, _, done, u, w = lax.while_loop(
-    cond_fun=cond_fn,
-    body_fun=body_fn,
-    init_val=(jnp.array(0), key, init_done, init_u, init_w),
-  )
-
-  return jnp.sign(u) * jnp.arccos(w)
-
-
-def _loc_scale(loc, scale, value):
-  if loc is None:
-    if scale is None:
-      return value
-    else:
-      return value * scale
-  else:
-    if scale is None:
-      return value + loc
-    else:
-      return value * scale + loc
-
-
-def _check_py_seq(seq):
-  return jnp.asarray(seq) if isinstance(seq, (tuple, list)) else seq
-
-
-# ---------------------------------------------------------------------------------------------------------------
-
-
 for __k in dir(RandomState):
   __t = getattr(RandomState, __k)
   if not __k.startswith('__') and callable(__t) and (not __t.__doc__):