brainstate 0.0.1__py2.py3-none-any.whl

brainstate/nn/_rate_rnns.py

@@ -0,0 +1,393 @@
+# 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 ..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 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 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 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 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 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

@@ -0,0 +1,130 @@
+# 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 ..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 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))
+
+  @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

@@ -0,0 +1,166 @@
+# 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 ..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

brainstate/nn/functional/__init__.py

@@ -0,0 +1,25 @@
+# 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 ._activations import *
+from ._activations import __all__ as __activations_all__
+from ._normalization import *
+from ._normalization import __all__ as __others_all__
+from ._spikes import *
+from ._spikes import __all__ as __spikes_all__
+
+__all__ = __spikes_all__ + __others_all__ + __activations_all__
+