brainstate 0.0.2.post20241009__py2.py3-none-any.whl → 0.1.0__py2.py3-none-any.whl

brainstate/nn/_embedding.py

@@ -1,66 +0,0 @@
-# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-from typing import Optional, Callable, Union
-
-from ._base import DnnLayer
-from .. import init
-from brainstate._state import ParamState
-from brainstate.mixin import Mode, Training
-from brainstate.typing import ArrayLike
-
-__all__ = [
-  'Embedding',
-]
-
-
-class Embedding(DnnLayer):
-  r"""
-  A simple lookup table that stores embeddings of a fixed size.
-
-  Args:
-    num_embeddings: Size of embedding dictionary. Must be non-negative.
-    embedding_size: Size of each embedding vector. Must be non-negative.
-    embed_init: The initializer for the embedding lookup table, of shape `(num_embeddings, embedding_size)`.
-
-  """
-
-  def __init__(
-      self,
-      num_embeddings: int,
-      embedding_size: int,
-      embed_init: Union[Callable, ArrayLike] = init.LecunUniform(),
-      name: Optional[str] = None,
-      mode: Optional[Mode] = None,
-  ):
-    super().__init__(name=name, mode=mode)
-    if num_embeddings < 0:
-      raise ValueError("num_embeddings must not be negative.")
-    if embedding_size < 0:
-      raise ValueError("embedding_size must not be negative.")
-    self.num_embeddings = num_embeddings
-    self.embedding_size = embedding_size
-    self.out_size = (embedding_size,)
-
-    weight = init.param(embed_init, (self.num_embeddings, self.embedding_size))
-    if self.mode.has(Training):
-      self.weight = ParamState(weight)
-    else:
-      self.weight = weight
-
-  def update(self, indices: ArrayLike):
-    if self.mode.has(Training):
-      return self.weight.value[indices]
-    return self.weight[indices]

brainstate/nn/_misc.py

@@ -1,133 +0,0 @@
-# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-
-from __future__ import annotations
-
-from enum import Enum
-from functools import wraps
-from typing import Sequence, Callable
-
-import brainunit as bu
-import jax.numpy as jnp
-
-from .. import environ
-from .._state import State
-from ..transform import vector_grad
-
-__all__ = [
-  # 'exp_euler',
-  'exp_euler_step',
-]
-
-git_issue_addr = 'https://github.com/brainpy/brainscale/issues'
-
-
-def state_traceback(states: Sequence[State]):
-  """
-  Traceback the states of the brain model.
-
-  Parameters
-  ----------
-  states : Sequence[bst.State]
-    The states of the brain model.
-
-  Returns
-  -------
-  str
-    The traceback information of the states.
-  """
-  state_info = []
-  for i, state in enumerate(states):
-    state_info.append(f'State {i}: {state}\n'
-                      f'defined at \n'
-                      f'{state.source_info.traceback}\n')
-  return '\n'.join(state_info)
-
-
-class BaseEnum(Enum):
-  @classmethod
-  def get_by_name(cls, name: str):
-    for item in cls:
-      if item.name == name:
-        return item
-    raise ValueError(f'Cannot find the {cls.__name__} type {name}.')
-
-  @classmethod
-  def get(cls, type_: str | Enum):
-    if isinstance(type_, cls):
-      return type_
-    elif isinstance(type_, str):
-      return cls.get_by_name(type_)
-    else:
-      raise ValueError(f'Cannot find the {cls.__name__} type {type_}.')
-
-
-def exp_euler(fun):
-  """
-  Exponential Euler method for solving ODEs.
-
-  Args:
-    fun: Callable. The function to be solved.
-
-  Returns:
-    The integral function.
-  """
-
-  @wraps(fun)
-  def integral(*args, **kwargs):
-    assert len(args) > 0, 'The input arguments should not be empty.'
-    if args[0].dtype not in [jnp.float32, jnp.float64, jnp.float16, jnp.bfloat16]:
-      raise ValueError(
-        'The input data type should be float32, float64, float16, or bfloat16 '
-        'when using Exponential Euler method.'
-        f'But we got {args[0].dtype}.'
-      )
-    dt = environ.get('dt')
-    linear, derivative = vector_grad(fun, argnums=0, return_value=True)(*args, **kwargs)
-    phi = bu.math.exprel(dt * linear)
-    return args[0] + dt * phi * derivative
-
-  return integral
-
-
-def exp_euler_step(fun: Callable, *args, **kwargs):
-  """
-  Exponential Euler method for solving ODEs.
-
-  Examples
-  --------
-  >>> def fun(x, t):
-  ...     return -x
-  >>> x = 1.0
-  >>> exp_euler_step(fun, x, None)
-
-  Args:
-    fun: Callable. The function to be solved.
-
-  Returns:
-    The integral function.
-  """
-  assert len(args) > 0, 'The input arguments should not be empty.'
-  if args[0].dtype not in [jnp.float32, jnp.float64, jnp.float16, jnp.bfloat16]:
-    raise ValueError(
-      'The input data type should be float32, float64, float16, or bfloat16 '
-      'when using Exponential Euler method.'
-      f'But we got {args[0].dtype}.'
-    )
-  dt = environ.get('dt')
-  linear, derivative = vector_grad(fun, argnums=0, return_value=True)(*args, **kwargs)
-  phi = bu.math.exprel(dt * linear)
-  return args[0] + dt * phi * derivative

brainstate/nn/_normalizations.py

@@ -1,389 +0,0 @@
-# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-# -*- coding: utf-8 -*-
-
-from __future__ import annotations
-
-import numbers
-from typing import Callable, Union, Sequence, Optional, Any
-
-import jax
-import jax.numpy as jnp
-
-from ._base import DnnLayer
-from .. import environ, init
-from brainstate._state import LongTermState, ParamState
-from brainstate.mixin import Mode
-from brainstate.typing import DTypeLike, ArrayLike, Size, Axes
-
-__all__ = [
-  'BatchNorm1d', 'BatchNorm2d', 'BatchNorm3d',
-]
-
-
-def _canonicalize_axes(ndim: int, feature_axes: Sequence[int]):
-  axes = []
-  for axis in feature_axes:
-    if axis < 0:
-      axis += ndim
-    if axis < 0 or axis >= ndim:
-      raise ValueError(f'Invalid axis {axis} for {ndim}D input')
-    axes.append(axis)
-  return tuple(axes)
-
-
-def _abs_sq(x):
-  """Computes the elementwise square of the absolute value |x|^2."""
-  if jnp.iscomplexobj(x):
-    return jax.lax.square(jax.lax.real(x)) + jax.lax.square(jax.lax.imag(x))
-  else:
-    return jax.lax.square(x)
-
-
-def _compute_stats(
-    x: ArrayLike,
-    axes: Sequence[int],
-    dtype: DTypeLike,
-    axis_name: Optional[str] = None,
-    axis_index_groups: Optional[Sequence[int]] = None,
-    use_mean: bool = True,
-):
-  """Computes mean and variance statistics.
-
-  This implementation takes care of a few important details:
-  - Computes in float32 precision for stability in half precision training.
-  - mean and variance are computable in a single XLA fusion,
-    by using Var = E[|x|^2] - |E[x]|^2 instead of Var = E[|x - E[x]|^2]).
-  - Clips negative variances to zero which can happen due to
-    roundoff errors. This avoids downstream NaNs.
-  - Supports averaging across a parallel axis and subgroups of a parallel axis
-    with a single `lax.pmean` call to avoid latency.
-
-  Arguments:
-    x: Input array.
-    axes: The axes in ``x`` to compute mean and variance statistics for.
-    dtype: tp.Optional dtype specifying the minimal precision. Statistics
-      are always at least float32 for stability (default: dtype of x).
-    axis_name: tp.Optional name for the pmapped axis to compute mean over.
-    axis_index_groups: tp.Optional axis indices.
-    use_mean: If true, calculate the mean from the input and use it when
-      computing the variance. If false, set the mean to zero and compute
-      the variance without subtracting the mean.
-
-  Returns:
-    A pair ``(mean, var)``.
-  """
-  if dtype is None:
-    dtype = jax.numpy.result_type(x)
-  # promote x to at least float32, this avoids half precision computation
-  # but preserves double or complex floating points
-  dtype = jax.numpy.promote_types(dtype, environ.dftype())
-  x = jnp.asarray(x, dtype)
-
-  # Compute mean and mean of squared values.
-  mean2 = jnp.mean(_abs_sq(x), axes)
-  if use_mean:
-    mean = jnp.mean(x, axes)
-  else:
-    mean = jnp.zeros(mean2.shape, dtype=dtype)
-
-  # If axis_name is provided, we need to average the mean and mean2 across
-  if axis_name is not None:
-    concatenated_mean = jnp.concatenate([mean, mean2])
-    mean, mean2 = jnp.split(
-      jax.lax.pmean(
-        concatenated_mean,
-        axis_name=axis_name,
-        axis_index_groups=axis_index_groups,
-      ),
-      2,
-    )
-
-  # mean2 - _abs_sq(mean) is not guaranteed to be non-negative due
-  # to floating point round-off errors.
-  var = jnp.maximum(0.0, mean2 - _abs_sq(mean))
-  return mean, var
-
-
-def _normalize(
-    x: ArrayLike,
-    mean: Optional[ArrayLike],
-    var: Optional[ArrayLike],
-    weights: Optional[ParamState],
-    reduction_axes: Sequence[int],
-    dtype: DTypeLike,
-    epsilon: Union[numbers.Number, jax.Array],
-):
-  """Normalizes the input of a normalization layer and optionally applies a learned scale and bias.
-
-  Arguments:
-    x: The input.
-    mean: Mean to use for normalization.
-    var: Variance to use for normalization.
-    weights: The scale and bias parameters.
-    reduction_axes: The axes in ``x`` to reduce.
-    dtype: The dtype of the result (default: infer from input and params).
-    epsilon: Normalization epsilon.
-
-  Returns:
-    The normalized input.
-  """
-  if mean is not None:
-    assert var is not None, 'mean and var must be both None or not None.'
-    stats_shape = list(x.shape)
-    for axis in reduction_axes:
-      stats_shape[axis] = 1
-    mean = mean.reshape(stats_shape)
-    var = var.reshape(stats_shape)
-    y = x - mean
-    mul = jax.lax.rsqrt(var + jnp.asarray(epsilon, dtype))
-    y = y * mul
-    if weights is not None:
-      y = _scale_operation(y, weights.value)
-  else:
-    assert var is None, 'mean and var must be both None or not None.'
-    assert weights is None, 'scale and bias are not supported without mean and var'
-    y = x
-  return jnp.asarray(y, dtype)
-
-
-def _scale_operation(x, param):
-  if 'scale' in param:
-    x = x * param['scale']
-  if 'bias' in param:
-    x = x + param['bias']
-  return x
-
-
-class _BatchNorm(DnnLayer):
-  __module__ = 'brainstate.nn'
-  num_spatial_dims: int
-
-  def __init__(
-      self,
-      in_size: Size,
-      feature_axis: Axes = -1,
-      track_running_stats: bool = True,
-      epsilon: float = 1e-5,
-      momentum: float = 0.99,
-      affine: bool = True,
-      bias_initializer: Union[ArrayLike, Callable] = init.Constant(0.),
-      scale_initializer: Union[ArrayLike, Callable] = init.Constant(1.),
-      axis_name: Optional[Union[str, Sequence[str]]] = None,
-      axis_index_groups: Optional[Sequence[Sequence[int]]] = None,
-      mode: Optional[Mode] = None,
-      name: Optional[str] = None,
-      dtype: Any = None,
-  ):
-    super().__init__(name=name, mode=mode)
-
-    # parameters
-    self.in_size = tuple(in_size)
-    self.out_size = tuple(in_size)
-    self.affine = affine
-    self.bias_initializer = bias_initializer
-    self.scale_initializer = scale_initializer
-    self.dtype = dtype or environ.dftype()
-    self.track_running_stats = track_running_stats
-    self.momentum = jnp.asarray(momentum, dtype=self.dtype)
-    self.epsilon = jnp.asarray(epsilon, dtype=self.dtype)
-
-    # parameters about axis
-    feature_axis = (feature_axis,) if isinstance(feature_axis, int) else feature_axis
-    self.feature_axis = _canonicalize_axes(len(in_size), feature_axis)
-    self.axis_name = axis_name
-    self.axis_index_groups = axis_index_groups
-
-    # variables
-    feature_shape = tuple([ax if i in self.feature_axis else 1 for i, ax in enumerate(in_size)])
-    if self.track_running_stats:
-      self.running_mean = LongTermState(jnp.zeros(feature_shape, dtype=self.dtype))
-      self.running_var = LongTermState(jnp.ones(feature_shape, dtype=self.dtype))
-    else:
-      self.running_mean = None
-      self.running_var = None
-
-    # parameters
-    if self.affine:
-      assert track_running_stats, "Affine parameters are not needed when track_running_stats is False."
-      bias = init.param(self.bias_initializer, feature_shape)
-      scale = init.param(self.scale_initializer, feature_shape)
-      self.weight = ParamState(dict(bias=bias, scale=scale))
-    else:
-      self.weight = None
-
-  def _check_input_dim(self, x):
-    if x.ndim == self.num_spatial_dims + 2:
-      x_shape = x.shape[1:]
-    elif x.ndim == self.num_spatial_dims + 1:
-      x_shape = x.shape
-    else:
-      raise ValueError(f"expected {self.num_spatial_dims + 2}D (with batch) or "
-                       f"{self.num_spatial_dims + 1}D (without batch) input (got {x.ndim}D input, {x.shape})")
-    if self.in_size != x_shape:
-      raise ValueError(f"The expected input shape is {self.in_size}, while we got {x_shape}.")
-
-  def update(self, x):
-    # input shape and batch mode or not
-    if x.ndim == self.num_spatial_dims + 2:
-      x_shape = x.shape[1:]
-      batch = True
-    elif x.ndim == self.num_spatial_dims + 1:
-      x_shape = x.shape
-      batch = False
-    else:
-      raise ValueError(f"expected {self.num_spatial_dims + 2}D (with batch) or "
-                       f"{self.num_spatial_dims + 1}D (without batch) input (got {x.ndim}D input, {x.shape})")
-    if self.in_size != x_shape:
-      raise ValueError(f"The expected input shape is {self.in_size}, while we got {x_shape}.")
-
-    # reduce the feature axis
-    if batch:
-      reduction_axes = tuple(i for i in range(x.ndim) if (i - 1) not in self.feature_axis)
-    else:
-      reduction_axes = tuple(i for i in range(x.ndim) if i not in self.feature_axis)
-
-    # fitting phase
-    fit_phase = environ.get('fit', desc='Whether this is a fitting process. Bool.')
-
-    # compute the running mean and variance
-    if self.track_running_stats:
-      if fit_phase:
-        mean, var = _compute_stats(
-          x,
-          reduction_axes,
-          dtype=self.dtype,
-          axis_name=self.axis_name,
-          axis_index_groups=self.axis_index_groups,
-        )
-        self.running_mean.value = self.momentum * self.running_mean.value + (1 - self.momentum) * mean
-        self.running_var.value = self.momentum * self.running_var.value + (1 - self.momentum) * var
-      else:
-        mean = self.running_mean.value
-        var = self.running_var.value
-    else:
-      mean, var = None, None
-
-    # normalize
-    return _normalize(x, mean, var, self.weight, reduction_axes, self.dtype, self.epsilon)
-
-
-class BatchNorm1d(_BatchNorm):
-  r"""1-D batch normalization [1]_.
-
-  The data should be of `(b, l, c)`, where `b` is the batch dimension,
-  `l` is the layer dimension, and `c` is the channel dimension.
-
-  %s
-  """
-  __module__ = 'brainstate.nn'
-  num_spatial_dims: int = 1
-
-
-class BatchNorm2d(_BatchNorm):
-  r"""2-D batch normalization [1]_.
-
-  The data should be of `(b, h, w, c)`, where `b` is the batch dimension,
-  `h` is the height dimension, `w` is the width dimension, and `c` is the
-  channel dimension.
-
-  %s
-  """
-  __module__ = 'brainstate.nn'
-  num_spatial_dims: int = 2
-
-
-class BatchNorm3d(_BatchNorm):
-  r"""3-D batch normalization [1]_.
-
-  The data should be of `(b, h, w, d, c)`, where `b` is the batch dimension,
-  `h` is the height dimension, `w` is the width dimension, `d` is the depth
-  dimension, and `c` is the channel dimension.
-
-  %s
-  """
-  __module__ = 'brainstate.nn'
-  num_spatial_dims: int = 3
-
-
-_bn_doc = r'''
-
-  This layer aims to reduce the internal covariant shift of data. It
-  normalizes a batch of data by fixing the mean and variance of inputs
-  on each feature (channel). Most commonly, the first axis of the data
-  is the batch, and the last is the channel. However, users can specify
-  the axes to be normalized.
-
-  .. math::
-     y=\frac{x-\mathrm{E}[x]}{\sqrt{\operatorname{Var}[x]+\epsilon}} * \gamma+\beta
-
-  .. note::
-      This :attr:`momentum` argument is different from one used in optimizer
-      classes and the conventional notion of momentum. Mathematically, the
-      update rule for running statistics here is
-      :math:`\hat{x}_\text{new} = \text{momentum} \times \hat{x} + (1-\text{momentum}) \times x_t`,
-      where :math:`\hat{x}` is the estimated statistic and :math:`x_t` is the
-      new observed value.
-
-  Parameters
-  ----------
-  in_size: sequence of int
-    The input shape, without batch size.
-  feature_axis: int, tuple, list
-    The feature or non-batch axis of the input.
-  track_running_stats: bool
-    A boolean value that when set to ``True``, this module tracks the running mean and variance, 
-    and when set to ``False``, this module does not track such statistics, and initializes 
-    statistics buffers ``running_mean`` and ``running_var`` as ``None``. When these buffers are ``None``, 
-    this module always uses batch statistics. in both training and eval modes. Default: ``True``.
-  momentum: float
-    The value used for the ``running_mean`` and ``running_var`` computation. Default: 0.99
-  epsilon: float
-    A value added to the denominator for numerical stability. Default: 1e-5
-  affine: bool
-    A boolean value that when set to ``True``, this module has
-    learnable affine parameters. Default: ``True``
-  bias_initializer: ArrayLike, Callable
-    An initializer generating the original translation matrix. If not ``None``, bias (beta) is added. 
-    Default: ``init.Constant(0.)``
-  scale_initializer: ArrayLike, Callable
-    An initializer generating the original scaling matrix. If not ``None``, multiply by scale (gamma).
-    Default: ``init.Constant(1.)``
-  axis_name: optional, str, sequence of str
-    If not ``None``, it should be a string (or sequence of
-    strings) representing the axis name(s) over which this module is being
-    run within a jax map (e.g. ``jax.pmap`` or ``jax.vmap``). Supplying this
-    argument means that batch statistics are calculated across all replicas
-    on the named axes.
-  axis_index_groups: optional, sequence
-    Specifies how devices are grouped. Valid
-    only within ``jax.pmap`` collectives.
-    Groups of axis indices within that named axis
-    representing subsets of devices to reduce over (default: None). For
-    example, `[[0, 1], [2, 3]]` would independently batch-normalize over
-    the examples on the first two and last two devices. See `jax.lax.psum`
-    for more details.
-
-  References
-  ----------
-  .. [1] Ioffe, Sergey and Christian Szegedy. “Batch Normalization: Accelerating Deep Network Training
-         by Reducing Internal Covariate Shift.” ArXiv abs/1502.03167 (2015): n. pag.
-
-'''
-
-BatchNorm1d.__doc__ = BatchNorm1d.__doc__ % _bn_doc
-BatchNorm2d.__doc__ = BatchNorm2d.__doc__ % _bn_doc
-BatchNorm3d.__doc__ = BatchNorm3d.__doc__ % _bn_doc

brainstate/nn/_others.py

@@ -1,101 +0,0 @@
-# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-
-from __future__ import annotations
-
-from functools import partial
-from typing import Optional
-
-import brainunit as bu
-import jax.numpy as jnp
-
-from ._base import DnnLayer
-from brainstate.mixin import Mode
-from brainstate import random, environ, typing, init
-
-__all__ = [
-  'DropoutFixed',
-]
-
-
-class DropoutFixed(DnnLayer):
-  """
-  A dropout layer with the fixed dropout mask along the time axis once after initialized.
-
-  In training, to compensate for the fraction of input values dropped (`rate`),
-  all surviving values are multiplied by `1 / (1 - rate)`.
-
-  This layer is active only during training (``mode=brainstate.mixin.Training``). In other
-  circumstances it is a no-op.
-
-  .. [1] Srivastava, Nitish, et al. "Dropout: a simple way to prevent
-         neural networks from overfitting." The journal of machine learning
-         research 15.1 (2014): 1929-1958.
-
-  .. admonition:: Tip
-      :class: tip
-
-      This kind of Dropout is firstly described in `Enabling Spike-based Backpropagation for Training Deep Neural
-      Network Architectures <https://arxiv.org/abs/1903.06379>`_:
-
-      There is a subtle difference in the way dropout is applied in SNNs compared to ANNs. In ANNs, each epoch of
-      training has several iterations of mini-batches. In each iteration, randomly selected units (with dropout ratio of :math:`p`)
-      are disconnected from the network while weighting by its posterior probability (:math:`1-p`). However, in SNNs, each
-      iteration has more than one forward propagation depending on the time length of the spike train. We back-propagate
-      the output error and modify the network parameters only at the last time step. For dropout to be effective in
-      our training method, it has to be ensured that the set of connected units within an iteration of mini-batch
-      data is not changed, such that the neural network is constituted by the same random subset of units during
-      each forward propagation within a single iteration. On the other hand, if the units are randomly connected at
-      each time-step, the effect of dropout will be averaged out over the entire forward propagation time within an
-      iteration. Then, the dropout effect would fade-out once the output error is propagated backward and the parameters
-      are updated at the last time step. Therefore, we need to keep the set of randomly connected units for the entire
-      time window within an iteration.
-
-  Args:
-    in_size: The size of the input tensor.
-    prob: Probability to keep element of the tensor.
-    mode: Mode. The computation mode of the object.
-    name: str. The name of the dynamic system.
-  """
-  __module__ = 'brainstate.nn'
-
-  def __init__(
-      self,
-      in_size: typing.Size,
-      prob: float = 0.5,
-      mode: Optional[Mode] = None,
-      name: Optional[str] = None
-  ) -> None:
-    super().__init__(mode=mode, name=name)
-    assert 0. <= prob < 1., f"Dropout probability must be in the range [0, 1). But got {prob}."
-    self.prob = prob
-    self.in_size = in_size
-    self.out_size = in_size
-
-  def init_state(self, batch_size=None, **kwargs):
-    self.mask = init.param(partial(random.bernoulli, self.prob), self.in_size, batch_size)
-
-  def update(self, x):
-    dtype = bu.math.get_dtype(x)
-    fit_phase = environ.get('fit', desc='Whether this is a fitting process. Bool.')
-    if fit_phase:
-      assert self.mask.shape == x.shape, (f"Input shape {x.shape} does not match the mask shape {self.mask.shape}. "
-                                          f"Please call `init_state()` method first.")
-      return jnp.where(self.mask,
-                       jnp.asarray(x / self.prob, dtype=dtype),
-                       jnp.asarray(0., dtype=dtype))
-    else:
-      return x