brainstate 0.0.2.post20241010__py2.py3-none-any.whl → 0.1.0.post20241122__py2.py3-none-any.whl

brainstate/nn/_rate_rnns.py

@@ -1,410 +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
-
-from typing import Callable, Union
-
-import jax.numpy as jnp
-
-from ._base import ExplicitInOutSize
-from ._connections import Linear
-from .. import random, init, functional
-from .._module import Module
-from .._state import ShortTermState, ParamState
-from ..mixin import DelayedInit, Mode
-from brainstate.typing import ArrayLike
-
-__all__ = [
-  'RNNCell', 'ValinaRNNCell', 'GRUCell', 'MGUCell', 'LSTMCell', 'URLSTMCell',
-]
-
-
-class RNNCell(Module, ExplicitInOutSize, DelayedInit):
-  """
-  Base class for RNN cells.
-  """
-  pass
-
-
-class ValinaRNNCell(RNNCell):
-  """
-  Vanilla RNN cell.
-
-  Args:
-    num_in: int. The number of input units.
-    num_out: int. The number of hidden units.
-    state_init: callable, ArrayLike. The state initializer.
-    w_init: callable, ArrayLike. The input weight initializer.
-    b_init: optional, callable, ArrayLike. The bias weight initializer.
-    activation: str, callable. The activation function. It can be a string or a callable function.
-    mode: optional, Mode. The mode of the module.
-    name: optional, str. The name of the module.
-  """
-  __module__ = 'brainstate.nn'
-
-  def __init__(
-      self,
-      num_in: int,
-      num_out: int,
-      state_init: Union[ArrayLike, Callable] = init.ZeroInit(),
-      w_init: Union[ArrayLike, Callable] = init.XavierNormal(),
-      b_init: Union[ArrayLike, Callable] = init.ZeroInit(),
-      activation: str | Callable = 'relu',
-      mode: Mode = None,
-      name: str = None,
-  ):
-    super().__init__(mode=mode, name=name)
-
-    # parameters
-    self._state_initializer = state_init
-    self.num_out = num_out
-    self.num_in = num_in
-    self.in_size = (num_in,)
-    self.out_size = (num_out,)
-
-    # activation function
-    if isinstance(activation, str):
-      self.activation = getattr(functional, activation)
-    else:
-      assert callable(activation), "The activation function should be a string or a callable function. "
-      self.activation = activation
-
-    # weights
-    self.W = Linear(num_in + num_out, num_out, w_init=w_init, b_init=b_init, name=self.name + '_W')
-
-  def init_state(self, batch_size: int = None, **kwargs):
-    self.h = ShortTermState(init.param(self._state_initializer, self.num_out, batch_size))
-
-  def reset_state(self, batch_size: int = None, **kwargs):
-    self.h.value = init.param(self._state_initializer, self.num_out, batch_size)
-
-  def update(self, x):
-    xh = jnp.concatenate([x, self.h.value], axis=-1)
-    h = self.W(xh)
-    self.h.value = self.activation(h)
-    return self.h.value
-
-
-class GRUCell(RNNCell):
-  """
-  Gated Recurrent Unit (GRU) cell.
-
-  Args:
-    num_in: int. The number of input units.
-    num_out: int. The number of hidden units.
-    state_init: callable, ArrayLike. The state initializer.
-    w_init: callable, ArrayLike. The input weight initializer.
-    b_init: optional, callable, ArrayLike. The bias weight initializer.
-    activation: str, callable. The activation function. It can be a string or a callable function.
-    mode: optional, Mode. The mode of the module.
-    name: optional, str. The name of the module.
-  """
-  __module__ = 'brainstate.nn'
-
-  def __init__(
-      self,
-      num_in: int,
-      num_out: int,
-      w_init: Union[ArrayLike, Callable] = init.Orthogonal(),
-      b_init: Union[ArrayLike, Callable] = init.ZeroInit(),
-      state_init: Union[ArrayLike, Callable] = init.ZeroInit(),
-      activation: str | Callable = 'tanh',
-      mode: Mode = None,
-      name: str = None,
-  ):
-    super().__init__(mode=mode, name=name)
-
-    # parameters
-    self._state_initializer = state_init
-    self.num_out = num_out
-    self.num_in = num_in
-    self.in_size = (num_in,)
-    self.out_size = (num_out,)
-
-    # activation function
-    if isinstance(activation, str):
-      self.activation = getattr(functional, activation)
-    else:
-      assert callable(activation), "The activation function should be a string or a callable function. "
-      self.activation = activation
-
-    # weights
-    self.Wrz = Linear(num_in + num_out, num_out * 2, w_init=w_init, b_init=b_init, name=self.name + '_Wrz')
-    self.Wh = Linear(num_in + num_out, num_out, w_init=w_init, b_init=b_init, name=self.name + '_Wh')
-
-  def init_state(self, batch_size: int = None, **kwargs):
-    self.h = ShortTermState(init.param(self._state_initializer, [self.num_out], batch_size))
-
-  def reset_state(self, batch_size: int = None, **kwargs):
-    self.h.value = init.param(self._state_initializer, [self.num_out], batch_size)
-
-  def update(self, x):
-    old_h = self.h.value
-    xh = jnp.concatenate([x, old_h], axis=-1)
-    r, z = jnp.split(functional.sigmoid(self.Wrz(xh)), indices_or_sections=2, axis=-1)
-    rh = r * old_h
-    h = self.activation(self.Wh(jnp.concatenate([x, rh], axis=-1)))
-    h = (1 - z) * old_h + z * h
-    self.h.value = h
-    return h
-
-
-class MGUCell(RNNCell):
-  r"""
-  Minimal Gated Recurrent Unit (MGU) cell.
-
-  .. math::
-
-     \begin{aligned}
-     f_{t}&=\sigma (W_{f}x_{t}+U_{f}h_{t-1}+b_{f})\\
-     {\hat {h}}_{t}&=\phi (W_{h}x_{t}+U_{h}(f_{t}\odot h_{t-1})+b_{h})\\
-     h_{t}&=(1-f_{t})\odot h_{t-1}+f_{t}\odot {\hat {h}}_{t}
-     \end{aligned}
-
-  where:
-
-  - :math:`x_{t}`: input vector
-  - :math:`h_{t}`: output vector
-  - :math:`{\hat {h}}_{t}`: candidate activation vector
-  - :math:`f_{t}`: forget vector
-  - :math:`W, U, b`: parameter matrices and vector
-
-  Args:
-    num_in: int. The number of input units.
-    num_out: int. The number of hidden units.
-    state_init: callable, ArrayLike. The state initializer.
-    w_init: callable, ArrayLike. The input weight initializer.
-    b_init: optional, callable, ArrayLike. The bias weight initializer.
-    activation: str, callable. The activation function. It can be a string or a callable function.
-    mode: optional, Mode. The mode of the module.
-    name: optional, str. The name of the module.
-  """
-  __module__ = 'brainstate.nn'
-
-  def __init__(
-      self,
-      num_in: int,
-      num_out: int,
-      w_init: Union[ArrayLike, Callable] = init.Orthogonal(),
-      b_init: Union[ArrayLike, Callable] = init.ZeroInit(),
-      state_init: Union[ArrayLike, Callable] = init.ZeroInit(),
-      activation: str | Callable = 'tanh',
-      mode: Mode = None,
-      name: str = None,
-  ):
-    super().__init__(mode=mode, name=name)
-
-    # parameters
-    self._state_initializer = state_init
-    self.num_out = num_out
-    self.num_in = num_in
-    self.in_size = (num_in,)
-    self.out_size = (num_out,)
-
-    # activation function
-    if isinstance(activation, str):
-      self.activation = getattr(functional, activation)
-    else:
-      assert callable(activation), "The activation function should be a string or a callable function. "
-      self.activation = activation
-
-    # weights
-    self.Wf = Linear(num_in + num_out, num_out, w_init=w_init, b_init=b_init, name=self.name + '_Wf')
-    self.Wh = Linear(num_in + num_out, num_out, w_init=w_init, b_init=b_init, name=self.name + '_Wh')
-
-  def init_state(self, batch_size: int = None, **kwargs):
-    self.h = ShortTermState(init.param(self._state_initializer, [self.num_out], batch_size))
-
-  def reset_state(self, batch_size: int = None, **kwargs):
-    self.h.value = init.param(self._state_initializer, [self.num_out], batch_size)
-
-  def update(self, x):
-    old_h = self.h.value
-    xh = jnp.concatenate([x, old_h], axis=-1)
-    f = functional.sigmoid(self.Wf(xh))
-    fh = f * old_h
-    h = self.activation(self.Wh(jnp.concatenate([x, fh], axis=-1)))
-    self.h.value = (1 - f) * self.h.value + f * h
-    return self.h.value
-
-
-class LSTMCell(RNNCell):
-  r"""Long short-term memory (LSTM) RNN core.
-
-  The implementation is based on (zaremba, et al., 2014) [1]_. Given
-  :math:`x_t` and the previous state :math:`(h_{t-1}, c_{t-1})` the core
-  computes
-
-  .. math::
-
-     \begin{array}{ll}
-     i_t = \sigma(W_{ii} x_t + W_{hi} h_{t-1} + b_i) \\
-     f_t = \sigma(W_{if} x_t + W_{hf} h_{t-1} + b_f) \\
-     g_t = \tanh(W_{ig} x_t + W_{hg} h_{t-1} + b_g) \\
-     o_t = \sigma(W_{io} x_t + W_{ho} h_{t-1} + b_o) \\
-     c_t = f_t c_{t-1} + i_t g_t \\
-     h_t = o_t \tanh(c_t)
-     \end{array}
-
-  where :math:`i_t`, :math:`f_t`, :math:`o_t` are input, forget and
-  output gate activations, and :math:`g_t` is a vector of cell updates.
-
-  The output is equal to the new hidden, :math:`h_t`.
-
-  Notes
-  -----
-
-  Forget gate initialization: Following (Jozefowicz, et al., 2015) [2]_ we add 1.0
-  to :math:`b_f` after initialization in order to reduce the scale of forgetting in
-  the beginning of the training.
-
-
-  Parameters
-  ----------
-  num_in: int
-    The dimension of the input vector
-  num_out: int
-    The number of hidden unit in the node.
-  state_init: callable, ArrayLike
-    The state initializer.
-  w_init: callable, ArrayLike
-    The input weight initializer.
-  b_init: optional, callable, ArrayLike
-    The bias weight initializer.
-  activation: str, callable
-    The activation function. It can be a string or a callable function.
-
-  References
-  ----------
-
-  .. [1] Zaremba, Wojciech, Ilya Sutskever, and Oriol Vinyals. "Recurrent neural
-         network regularization." arXiv preprint arXiv:1409.2329 (2014).
-  .. [2] Jozefowicz, Rafal, Wojciech Zaremba, and Ilya Sutskever. "An empirical
-         exploration of recurrent network architectures." In International conference
-         on machine learning, pp. 2342-2350. PMLR, 2015.
-  """
-  __module__ = 'brainstate.nn'
-
-  def __init__(
-      self,
-      num_in: int,
-      num_out: int,
-      w_init: Union[ArrayLike, Callable] = init.XavierNormal(),
-      b_init: Union[ArrayLike, Callable] = init.ZeroInit(),
-      state_init: Union[ArrayLike, Callable] = init.ZeroInit(),
-      activation: str | Callable = 'tanh',
-      mode: Mode = None,
-      name: str = None,
-  ):
-    super().__init__(mode=mode, name=name)
-
-    # parameters
-    self.num_out = num_out
-    self.num_in = num_in
-    self.in_size = (num_in,)
-    self.out_size = (num_out,)
-
-    # initializers
-    self._state_initializer = state_init
-
-    # activation function
-    if isinstance(activation, str):
-      self.activation = getattr(functional, activation)
-    else:
-      assert callable(activation), "The activation function should be a string or a callable function. "
-      self.activation = activation
-
-    # weights
-    self.W = Linear(num_in + num_out, num_out * 4, w_init=w_init, b_init=b_init, name=self.name + '_W')
-
-  def init_state(self, batch_size: int = None, **kwargs):
-    self.c = ShortTermState(init.param(self._state_initializer, [self.num_out], batch_size))
-    self.h = ShortTermState(init.param(self._state_initializer, [self.num_out], batch_size))
-
-  def reset_state(self, batch_size: int = None, **kwargs):
-    self.c.value = init.param(self._state_initializer, [self.num_out], batch_size)
-    self.h.value = init.param(self._state_initializer, [self.num_out], batch_size)
-
-  def update(self, x):
-    h, c = self.h.value, self.c.value
-    xh = jnp.concat([x, h], axis=-1)
-    i, g, f, o = jnp.split(self.W(xh), indices_or_sections=4, axis=-1)
-    c = functional.sigmoid(f + 1.) * c + functional.sigmoid(i) * self.activation(g)
-    h = functional.sigmoid(o) * self.activation(c)
-    self.h.value = h
-    self.c.value = c
-    return h
-
-
-class URLSTMCell(RNNCell):
-  def __init__(
-      self,
-      num_in: int,
-      num_out: int,
-      w_init: Union[ArrayLike, Callable] = init.XavierNormal(),
-      state_init: Union[ArrayLike, Callable] = init.ZeroInit(),
-      activation: str | Callable = 'tanh',
-      mode: Mode = None,
-      name: str = None,
-  ):
-    super().__init__(mode=mode, name=name)
-
-    # parameters
-    self.num_out = num_out
-    self.num_in = num_in
-    self.in_size = (num_in,)
-    self.out_size = (num_out,)
-
-    # initializers
-    self._state_initializer = state_init
-
-    # activation function
-    if isinstance(activation, str):
-      self.activation = getattr(functional, activation)
-    else:
-      assert callable(activation), "The activation function should be a string or a callable function. "
-      self.activation = activation
-
-    # weights
-    self.W = Linear(num_in + num_out, num_out * 4, w_init=w_init, b_init=None, name=self.name + '_Wg')
-    self.bias = ParamState(self._forget_bias())
-
-  def _forget_bias(self):
-    u = random.uniform(1 / self.num_out, 1 - 1 / self.num_out, (self.num_out,))
-    return -jnp.log(1 / u - 1)
-
-  def init_state(self, batch_size: int = None, **kwargs):
-    self.c = ShortTermState(init.param(self._state_initializer, [self.num_out], batch_size))
-    self.h = ShortTermState(init.param(self._state_initializer, [self.num_out], batch_size))
-
-  def reset_state(self, batch_size: int = None, **kwargs):
-    self.c.value = init.param(self._state_initializer, [self.num_out], batch_size)
-    self.h.value = init.param(self._state_initializer, [self.num_out], batch_size)
-
-  def update(self, x: ArrayLike) -> ArrayLike:
-    h, c = self.h.value, self.c.value
-    xh = jnp.concat([x, h], axis=-1)
-    f, r, u, o = jnp.split(self.W(xh), indices_or_sections=4, axis=-1)
-    f_ = functional.sigmoid(f + self.bias.value)
-    r_ = functional.sigmoid(r - self.bias.value)
-    g = 2 * r_ * f_ + (1 - 2 * r_) * f_ ** 2
-    next_cell = g * c + (1 - g) * self.activation(u)
-    next_hidden = functional.sigmoid(o) * self.activation(next_cell)
-    self.h.value = next_hidden
-    self.c.value = next_cell
-    return next_hidden

brainstate/nn/_readout.py

@@ -1,136 +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
-
-import jax
-import jax.numpy as jnp
-
-from ._base import DnnLayer
-from ._dynamics import Neuron
-from ._misc import exp_euler_step
-from .. import environ, init, surrogate
-from .._state import ShortTermState, ParamState
-from ..mixin import Mode
-from brainstate.typing import Size, ArrayLike, DTypeLike
-
-__all__ = [
-  'LeakyRateReadout',
-  'LeakySpikeReadout',
-]
-
-
-class LeakyRateReadout(DnnLayer):
-  """
-  Leaky dynamics for the read-out module used in the Real-Time Recurrent Learning.
-  """
-  __module__ = 'brainstate.nn'
-
-  def __init__(
-      self,
-      in_size: Size,
-      out_size: Size,
-      tau: ArrayLike = 5.,
-      w_init: Callable = init.KaimingNormal(),
-      mode: Mode = None,
-      name: str = None,
-  ):
-    super().__init__(mode=mode, name=name)
-
-    # parameters
-    self.in_size = (in_size,) if isinstance(in_size, numbers.Integral) else tuple(in_size)
-    self.out_size = (out_size,) if isinstance(out_size, numbers.Integral) else tuple(out_size)
-    self.tau = init.param(tau, self.in_size)
-    self.decay = jnp.exp(-environ.get_dt() / self.tau)
-
-    # weights
-    self.weight = ParamState(init.param(w_init, (self.in_size[0], self.out_size[0])))
-
-  def init_state(self, batch_size=None, **kwargs):
-    self.r = ShortTermState(init.param(init.Constant(0.), self.out_size, batch_size))
-
-  def reset_state(self, batch_size=None, **kwargs):
-    self.r.value = init.param(init.Constant(0.), self.out_size, batch_size)
-
-  def update(self, x):
-    r = self.decay * self.r.value + x @ self.weight.value
-    self.r.value = r
-    return r
-
-
-class LeakySpikeReadout(Neuron):
-  """
-  Integrate-and-fire neuron model.
-  """
-
-  __module__ = 'brainstate.nn'
-
-  def __init__(
-      self,
-      in_size: Size,
-      keep_size: bool = False,
-      tau: ArrayLike = 5.,
-      V_th: ArrayLike = 1.,
-      w_init: Callable = init.KaimingNormal(),
-      spk_fun: Callable = surrogate.ReluGrad(),
-      spk_dtype: DTypeLike = None,
-      spk_reset: str = 'soft',
-      mode: Mode = None,
-      name: str = None,
-  ):
-    super().__init__(in_size, keep_size=keep_size, name=name, mode=mode,
-                     spk_fun=spk_fun, spk_dtype=spk_dtype, spk_reset=spk_reset)
-
-    # parameters
-    self.tau = init.param(tau, (self.num,))
-    self.V_th = init.param(V_th, (self.num,))
-
-    # weights
-    self.weight = ParamState(init.param(w_init, (self.in_size[0], self.out_size[0])))
-
-  def dv(self, v, t, x):
-    x = self.sum_current_inputs(v, init=x)
-    return (-v + x) / self.tau
-
-  def init_state(self, batch_size, **kwargs):
-    self.V = ShortTermState(init.param(init.Constant(0.), self.varshape, batch_size))
-
-  def reset_state(self, batch_size, **kwargs):
-    self.V.value = init.param(init.Constant(0.), self.varshape, batch_size)
-
-  @property
-  def spike(self):
-    return self.get_spike(self.V.value)
-
-  def get_spike(self, V):
-    v_scaled = (V - self.V_th) / self.V_th
-    return self.spk_fun(v_scaled)
-
-  def update(self, x):
-    # reset
-    last_V = self.V.value
-    last_spike = self.get_spike(last_V)
-    V_th = self.V_th if self.spk_reset == 'soft' else jax.lax.stop_gradient(last_V)
-    V = last_V - V_th * last_spike
-    # membrane potential
-    V = exp_euler_step(self.dv, V, environ.get('t'), x @ self.weight.value)
-    V = V + self.sum_delta_inputs()
-    self.V.value = V
-    return self.get_spike(V)

brainstate/nn/_synouts.py

@@ -1,166 +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
-
-from typing import Optional
-
-import jax.numpy as jnp
-
-from .._module import Module
-from ..mixin import DelayedInit, BindCondData
-from brainstate.typing import ArrayLike
-
-__all__ = [
-  'SynOut', 'COBA', 'CUBA', 'MgBlock',
-]
-
-
-class SynOut(Module, DelayedInit, BindCondData):
-  """
-  Base class for synaptic outputs.
-
-  :py:class:`~.SynOut` is also subclass of :py:class:`~.ParamDesc` and :py:class:`~.BindCondData`.
-  """
-
-  __module__ = 'brainstate.nn'
-
-  def __init__(self, name: Optional[str] = None):
-    super().__init__(name=name)
-    self._conductance = None
-
-  def __call__(self, *args, **kwargs):
-    if self._conductance is None:
-      raise ValueError(f'Please first pack conductance data at the current step using '
-                       f'".{BindCondData.bind_cond.__name__}(data)". {self}')
-    ret = self.update(self._conductance, *args, **kwargs)
-    return ret
-
-
-class COBA(SynOut):
-  r"""
-  Conductance-based synaptic output.
-
-  Given the synaptic conductance, the model output the post-synaptic current with
-
-  .. math::
-
-     I_{syn}(t) = g_{\mathrm{syn}}(t) (E - V(t))
-
-  Parameters
-  ----------
-  E: ArrayLike
-    The reversal potential.
-  name: str
-    The model name.
-
-  See Also
-  --------
-  CUBA
-  """
-  __module__ = 'brainstate.nn'
-
-  def __init__(self, E: ArrayLike, name: Optional[str] = None):
-    super().__init__(name=name)
-
-    self.E = E
-
-  def update(self, conductance, potential):
-    return conductance * (self.E - potential)
-
-
-class CUBA(SynOut):
-  r"""Current-based synaptic output.
-
-  Given the conductance, this model outputs the post-synaptic current with a identity function:
-
-  .. math::
-
-     I_{\mathrm{syn}}(t) = g_{\mathrm{syn}}(t)
-
-  Parameters
-  ----------
-  name: str
-    The model name.
-
-  See Also
-  --------
-  COBA
-  """
-  __module__ = 'brainstate.nn'
-
-  def __init__(self, name: Optional[str] = None, ):
-    super().__init__(name=name)
-
-  def update(self, conductance, potential=None):
-    return conductance
-
-
-class MgBlock(SynOut):
-  r"""Synaptic output based on Magnesium blocking.
-
-  Given the synaptic conductance, the model output the post-synaptic current with
-
-  .. math::
-
-     I_{syn}(t) = g_{\mathrm{syn}}(t) (E - V(t)) g_{\infty}(V,[{Mg}^{2+}]_{o})
-
-  where The fraction of channels :math:`g_{\infty}` that are not blocked by magnesium can be fitted to
-
-  .. math::
-
-     g_{\infty}(V,[{Mg}^{2+}]_{o}) = (1+{e}^{-\alpha V} \frac{[{Mg}^{2+}]_{o}} {\beta})^{-1}
-
-  Here :math:`[{Mg}^{2+}]_{o}` is the extracellular magnesium concentration.
-
-  Parameters
-  ----------
-  E: ArrayLike
-    The reversal potential for the synaptic current. [mV]
-  alpha: ArrayLike
-    Binding constant. Default 0.062
-  beta: ArrayLike
-    Unbinding constant. Default 3.57
-  cc_Mg: ArrayLike
-    Concentration of Magnesium ion. Default 1.2 [mM].
-  V_offset: ArrayLike
-    The offset potential. Default 0. [mV]
-  name: str
-    The model name.
-  """
-  __module__ = 'brainstate.nn'
-
-  def __init__(
-      self,
-      E: ArrayLike = 0.,
-      cc_Mg: ArrayLike = 1.2,
-      alpha: ArrayLike = 0.062,
-      beta: ArrayLike = 3.57,
-      V_offset: ArrayLike = 0.,
-      name: Optional[str] = None,
-  ):
-    super().__init__(name=name)
-
-    self.E = E
-    self.V_offset = V_offset
-    self.cc_Mg = cc_Mg
-    self.alpha = alpha
-    self.beta = beta
-
-  def update(self, conductance, potential):
-    norm = (1 + self.cc_Mg / self.beta * jnp.exp(self.alpha * (self.V_offset - potential)))
-    return conductance * (self.E - potential) / norm