brainstate 0.1.1__py2.py3-none-any.whl → 0.1.3__py2.py3-none-any.whl

brainstate/nn/{_dynamics/_projection_base.py → _projection.py}

@@ -17,10 +17,10 @@ from typing import Union, Callable, Optional
 
 from brainstate._state import State
 from brainstate.mixin import AlignPost, ParamDescriber, BindCondData, JointTypes
-from brainstate.nn._collective_ops import call_order
-from brainstate.nn._module import Module
 from brainstate.util._others import get_unique_name
-from ._dynamics_base import Dynamics, maybe_init_prefetch, Prefetch, PrefetchDelayAt
+from ._collective_ops import call_order
+from ._dynamics import Dynamics, maybe_init_prefetch, Prefetch, PrefetchDelayAt
+from ._module import Module
 from ._synouts import SynOut
 
 __all__ = [
@@ -360,3 +360,35 @@ class CurrentProj(Interaction):
         x = self.prefetch(*x)
         x = self.comm(x)
         self.out.bind_cond(x)
+
+
+class RawProj(Interaction):
+    """
+    """
+    __module__ = 'brainstate.nn'
+
+    def __init__(
+        self,
+        comm: Callable,
+        out: SynOut,
+        post: Dynamics,
+    ):
+        super().__init__(name=get_unique_name(self.__class__.__name__))
+
+        # check out
+        if not isinstance(out, SynOut):
+            raise TypeError(f'The out should be a SynOut, but got {out}.')
+        self.out = out
+
+        # check post
+        if not isinstance(post, Dynamics):
+            raise TypeError(f'The post should be a Dynamics, but got {post}.')
+        self.post = post
+        post.add_current_input(self.name, out)
+
+        # output initialization
+        self.comm = comm
+
+    def update(self, x):
+        x = self.comm(x)
+        self.out.bind_cond(x)

brainstate/nn/{_dyn_impl/_rate_rnns.py → _rate_rnns.py}

@@ -22,9 +22,9 @@ import jax.numpy as jnp
 
 from brainstate import random, init, functional
 from brainstate._state import HiddenState, ParamState
-from brainstate.nn._interaction._linear import Linear
-from brainstate.nn._module import Module
 from brainstate.typing import ArrayLike
+from ._linear import Linear
+from ._module import Module
 
 __all__ = [
     'RNNCell', 'ValinaRNNCell', 'GRUCell', 'MGUCell', 'LSTMCell', 'URLSTMCell',

brainstate/nn/{_dyn_impl/_rate_rnns_test.py → _rate_rnns_test.py}

@@ -13,13 +13,12 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
 
 import unittest
 
 import jax.numpy as jnp
 
-import brainstate as bst
+import brainstate
 
 
 class TestRateRNNModels(unittest.TestCase):
@@ -30,31 +29,31 @@ class TestRateRNNModels(unittest.TestCase):
         self.x = jnp.ones((self.batch_size, self.num_in))
 
     def test_ValinaRNNCell(self):
-        model = bst.nn.ValinaRNNCell(num_in=self.num_in, num_out=self.num_out)
+        model = brainstate.nn.ValinaRNNCell(num_in=self.num_in, num_out=self.num_out)
         model.init_state(batch_size=self.batch_size)
         output = model.update(self.x)
         self.assertEqual(output.shape, (self.batch_size, self.num_out))
 
     def test_GRUCell(self):
-        model = bst.nn.GRUCell(num_in=self.num_in, num_out=self.num_out)
+        model = brainstate.nn.GRUCell(num_in=self.num_in, num_out=self.num_out)
         model.init_state(batch_size=self.batch_size)
         output = model.update(self.x)
         self.assertEqual(output.shape, (self.batch_size, self.num_out))
 
     def test_MGUCell(self):
-        model = bst.nn.MGUCell(num_in=self.num_in, num_out=self.num_out)
+        model = brainstate.nn.MGUCell(num_in=self.num_in, num_out=self.num_out)
         model.init_state(batch_size=self.batch_size)
         output = model.update(self.x)
         self.assertEqual(output.shape, (self.batch_size, self.num_out))
 
     def test_LSTMCell(self):
-        model = bst.nn.LSTMCell(num_in=self.num_in, num_out=self.num_out)
+        model = brainstate.nn.LSTMCell(num_in=self.num_in, num_out=self.num_out)
         model.init_state(batch_size=self.batch_size)
         output = model.update(self.x)
         self.assertEqual(output.shape, (self.batch_size, self.num_out))
 
     def test_URLSTMCell(self):
-        model = bst.nn.URLSTMCell(num_in=self.num_in, num_out=self.num_out)
+        model = brainstate.nn.URLSTMCell(num_in=self.num_in, num_out=self.num_out)
         model.init_state(batch_size=self.batch_size)
         output = model.update(self.x)
         self.assertEqual(output.shape, (self.batch_size, self.num_out))

brainstate/nn/{_dyn_impl/_readout.py → _readout.py}

@@ -24,10 +24,10 @@ import jax
 
 from brainstate import environ, init, surrogate
 from brainstate._state import HiddenState, ParamState
-from brainstate.nn._exp_euler import exp_euler_step
-from brainstate.nn._module import Module
 from brainstate.typing import Size, ArrayLike
-from ._dynamics_neuron import Neuron
+from ._exp_euler import exp_euler_step
+from ._module import Module
+from ._neuron import Neuron
 
 __all__ = [
     'LeakyRateReadout',

brainstate/nn/{_dyn_impl/_readout_test.py → _readout_test.py}

@@ -13,13 +13,12 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
 
 import unittest
 
 import jax.numpy as jnp
 
-import brainstate as bst
+import brainstate
 
 
 class TestReadoutModels(unittest.TestCase):
@@ -32,23 +31,23 @@ class TestReadoutModels(unittest.TestCase):
         self.x = jnp.ones((self.batch_size, self.in_size))
 
     def test_LeakyRateReadout(self):
-        with bst.environ.context(dt=0.1):
-            model = bst.nn.LeakyRateReadout(in_size=self.in_size, out_size=self.out_size, tau=self.tau)
+        with brainstate.environ.context(dt=0.1):
+            model = brainstate.nn.LeakyRateReadout(in_size=self.in_size, out_size=self.out_size, tau=self.tau)
             model.init_state(batch_size=self.batch_size)
             output = model.update(self.x)
             self.assertEqual(output.shape, (self.batch_size, self.out_size))
 
     def test_LeakySpikeReadout(self):
-        with bst.environ.context(dt=0.1):
-            model = bst.nn.LeakySpikeReadout(in_size=self.in_size, tau=self.tau, V_th=self.V_th,
-                                             V_initializer=bst.init.ZeroInit(),
-                                             w_init=bst.init.KaimingNormal())
+        with brainstate.environ.context(dt=0.1):
+            model = brainstate.nn.LeakySpikeReadout(in_size=self.in_size, tau=self.tau, V_th=self.V_th,
+                                                    V_initializer=brainstate.init.ZeroInit(),
+                                                    w_init=brainstate.init.KaimingNormal())
             model.init_state(batch_size=self.batch_size)
-            with bst.environ.context(t=0.):
+            with brainstate.environ.context(t=0.):
                 output = model.update(self.x)
             self.assertEqual(output.shape, (self.batch_size, self.out_size))
 
 
 if __name__ == '__main__':
-    with bst.environ.context(dt=0.1):
+    with brainstate.environ.context(dt=0.1):
         unittest.main()

brainstate/nn/_stp.py

@@ -0,0 +1,236 @@
+# Copyright 2025 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 typing import Optional
+
+import brainunit as u
+
+from brainstate import init
+from brainstate._state import HiddenState
+from brainstate.typing import ArrayLike, Size
+from ._exp_euler import exp_euler_step
+from ._synapse import Synapse
+
+__all__ = [
+    'ShortTermPlasticity', 'STP', 'STD',
+]
+
+
+class ShortTermPlasticity(Synapse):
+    pass
+
+
+class STP(ShortTermPlasticity):
+    r"""
+    Synapse with short-term plasticity.
+
+    This class implements a synapse model with short-term plasticity (STP), which captures
+    activity-dependent changes in synaptic efficacy that occur over milliseconds to seconds.
+    The model simultaneously accounts for both short-term facilitation and depression
+    based on the formulation by Tsodyks & Markram (1998).
+
+    The model is characterized by the following equations:
+
+    $$
+    \frac{du}{dt} = -\frac{u}{\tau_f} + U \cdot (1 - u) \cdot \delta(t - t_{spike})
+    $$
+
+    $$
+    \frac{dx}{dt} = \frac{1 - x}{\tau_d} - u \cdot x \cdot \delta(t - t_{spike})
+    $$
+
+    $$
+    g_{syn} = u \cdot x
+    $$
+
+    where:
+    - $u$ represents the utilization of synaptic efficacy (facilitation variable)
+    - $x$ represents the available synaptic resources (depression variable)
+    - $\tau_f$ is the facilitation time constant
+    - $\tau_d$ is the depression time constant
+    - $U$ is the baseline utilization parameter
+    - $\delta(t - t_{spike})$ is the Dirac delta function representing presynaptic spikes
+    - $g_{syn}$ is the effective synaptic conductance
+
+    Parameters
+    ----------
+    in_size : Size
+        Size of the input.
+    name : str, optional
+        Name of the synapse instance.
+    U : ArrayLike, default=0.15
+        Baseline utilization parameter (fraction of resources used per action potential).
+    tau_f : ArrayLike, default=1500.*u.ms
+        Time constant of short-term facilitation in milliseconds.
+    tau_d : ArrayLike, default=200.*u.ms
+        Time constant of short-term depression (recovery of synaptic resources) in milliseconds.
+
+    Attributes
+    ----------
+    u : HiddenState
+        Utilization of synaptic efficacy (facilitation variable).
+    x : HiddenState
+        Available synaptic resources (depression variable).
+
+    Notes
+    -----
+    - Larger values of tau_f produce stronger facilitation effects.
+    - Larger values of tau_d lead to slower recovery from depression.
+    - The parameter U controls the initial release probability.
+    - The effective synaptic strength is the product of u and x.
+
+    References
+    ----------
+    .. [1] Tsodyks, M. V., & Markram, H. (1997). The neural code between neocortical
+           pyramidal neurons depends on neurotransmitter release probability.
+           Proceedings of the National Academy of Sciences, 94(2), 719-723.
+    .. [2] Tsodyks, M., Pawelzik, K., & Markram, H. (1998). Neural networks with dynamic
+           synapses. Neural computation, 10(4), 821-835.
+    """
+    __module__ = 'brainstate.nn'
+
+    def __init__(
+        self,
+        in_size: Size,
+        name: Optional[str] = None,
+        U: ArrayLike = 0.15,
+        tau_f: ArrayLike = 1500. * u.ms,
+        tau_d: ArrayLike = 200. * u.ms,
+    ):
+        super().__init__(name=name, in_size=in_size)
+
+        # parameters
+        self.tau_f = init.param(tau_f, self.varshape)
+        self.tau_d = init.param(tau_d, self.varshape)
+        self.U = init.param(U, self.varshape)
+
+    def init_state(self, batch_size: int = None, **kwargs):
+        self.x = HiddenState(init.param(init.Constant(1.), self.varshape, batch_size))
+        self.u = HiddenState(init.param(init.Constant(self.U), self.varshape, batch_size))
+
+    def reset_state(self, batch_size: int = None, **kwargs):
+        self.x.value = init.param(init.Constant(1.), self.varshape, batch_size)
+        self.u.value = init.param(init.Constant(self.U), self.varshape, batch_size)
+
+    def update(self, pre_spike):
+        u = exp_euler_step(lambda u: - u / self.tau_f, self.u.value)
+        x = exp_euler_step(lambda x: (1 - x) / self.tau_d, self.x.value)
+
+        # --- original code:
+        #   if pre_spike.dtype == jax.numpy.bool_:
+        #     u = bm.where(pre_spike, u + self.U * (1 - self.u), u)
+        #     x = bm.where(pre_spike, x - u * self.x, x)
+        #   else:
+        #     u = pre_spike * (u + self.U * (1 - self.u)) + (1 - pre_spike) * u
+        #     x = pre_spike * (x - u * self.x) + (1 - pre_spike) * x
+
+        # --- simplified code:
+        u = u + pre_spike * self.U * (1 - self.u.value)
+        x = x - pre_spike * u * self.x.value
+
+        self.u.value = u
+        self.x.value = x
+        return u * x * pre_spike
+
+
+class STD(ShortTermPlasticity):
+    r"""
+    Synapse with short-term depression.
+
+    This class implements a synapse model with short-term depression (STD), which captures
+    activity-dependent reduction in synaptic efficacy, typically caused by depletion of
+    neurotransmitter vesicles following repeated stimulation.
+
+    The model is characterized by the following equation:
+
+    $$
+    \frac{dx}{dt} = \frac{1 - x}{\tau} - U \cdot x \cdot \delta(t - t_{spike})
+    $$
+
+    $$
+    g_{syn} = x
+    $$
+
+    where:
+    - $x$ represents the available synaptic resources (depression variable)
+    - $\tau$ is the depression recovery time constant
+    - $U$ is the utilization parameter (fraction of resources depleted per spike)
+    - $\delta(t - t_{spike})$ is the Dirac delta function representing presynaptic spikes
+    - $g_{syn}$ is the effective synaptic conductance
+
+    Parameters
+    ----------
+    in_size : Size
+        Size of the input.
+    name : str, optional
+        Name of the synapse instance.
+    tau : ArrayLike, default=200.*u.ms
+        Time constant governing recovery of synaptic resources in milliseconds.
+    U : ArrayLike, default=0.07
+        Utilization parameter (fraction of resources used per action potential).
+
+    Attributes
+    ----------
+    x : HiddenState
+        Available synaptic resources (depression variable).
+
+    Notes
+    -----
+    - Larger values of tau lead to slower recovery from depression.
+    - Larger values of U cause stronger depression with each spike.
+    - This model is a simplified version of the STP model that only includes depression.
+
+    References
+    ----------
+    .. [1] Abbott, L. F., Varela, J. A., Sen, K., & Nelson, S. B. (1997). Synaptic
+           depression and cortical gain control. Science, 275(5297), 220-224.
+    .. [2] Tsodyks, M. V., & Markram, H. (1997). The neural code between neocortical
+           pyramidal neurons depends on neurotransmitter release probability.
+           Proceedings of the National Academy of Sciences, 94(2), 719-723.
+    """
+    __module__ = 'brainstate.nn'
+
+    def __init__(
+        self,
+        in_size: Size,
+        name: Optional[str] = None,
+        # synapse parameters
+        tau: ArrayLike = 200. * u.ms,
+        U: ArrayLike = 0.07,
+    ):
+        super().__init__(name=name, in_size=in_size)
+
+        # parameters
+        self.tau = init.param(tau, self.varshape)
+        self.U = init.param(U, self.varshape)
+
+    def init_state(self, batch_size: int = None, **kwargs):
+        self.x = HiddenState(init.param(init.Constant(1.), self.varshape, batch_size))
+
+    def reset_state(self, batch_size: int = None, **kwargs):
+        self.x.value = init.param(init.Constant(1.), self.varshape, batch_size)
+
+    def update(self, pre_spike):
+        x = exp_euler_step(lambda x: (1 - x) / self.tau, self.x.value)
+
+        # --- original code:
+        # self.x.value = bm.where(pre_spike, x - self.U * self.x, x)
+
+        # --- simplified code:
+        self.x.value = x - pre_spike * self.U * self.x.value
+
+        return self.x.value * pre_spike

brainstate/nn/{_dyn_impl/_dynamics_synapse.py → _synapse.py}

@@ -23,12 +23,12 @@ import brainunit as u
 from brainstate import init, environ
 from brainstate._state import ShortTermState, HiddenState
 from brainstate.mixin import AlignPost
-from brainstate.nn._dynamics._dynamics_base import Dynamics
-from brainstate.nn._exp_euler import exp_euler_step
-from brainstate.typing import ArrayLike, Size
+from brainstate.typing import ArrayLike, Size, PyTree
+from ._dynamics import Dynamics
+from ._exp_euler import exp_euler_step
 
 __all__ = [
-    'Synapse', 'Expon', 'DualExpon', 'Alpha', 'STP', 'STD', 'AMPA', 'GABAa',
+    'Synapse', 'Expon', 'DualExpon', 'Alpha', 'AMPA', 'GABAa',
 ]
 
 
@@ -123,6 +123,9 @@ class Expon(Synapse, AlignPost):
         g = exp_euler_step(lambda g: self.sum_current_inputs(-g) / self.tau, self.g.value)
         self.g.value = self.sum_delta_inputs(g)
         if x is not None: self.g.value += x
+        return self.update_return()
+
+    def update_return(self) -> PyTree:
         return self.g.value
 
 
@@ -229,6 +232,9 @@ class DualExpon(Synapse, AlignPost):
         if x is not None:
             self.g_rise.value += x
             self.g_decay.value += x
+        return self.update_return()
+
+    def update_return(self) -> PyTree:
         return self.a * (self.g_decay.value - self.g_rise.value)
 
 
@@ -301,209 +307,10 @@ class Alpha(Synapse):
         self.h.value = self.sum_delta_inputs(h)
         if x is not None:
             self.h.value += x
-        return self.g.value
-
-
-class STP(Synapse):
-    r"""
-    Synapse with short-term plasticity.
-
-    This class implements a synapse model with short-term plasticity (STP), which captures
-    activity-dependent changes in synaptic efficacy that occur over milliseconds to seconds.
-    The model simultaneously accounts for both short-term facilitation and depression
-    based on the formulation by Tsodyks & Markram (1998).
-
-    The model is characterized by the following equations:
-
-    $$
-    \frac{du}{dt} = -\frac{u}{\tau_f} + U \cdot (1 - u) \cdot \delta(t - t_{spike})
-    $$
-
-    $$
-    \frac{dx}{dt} = \frac{1 - x}{\tau_d} - u \cdot x \cdot \delta(t - t_{spike})
-    $$
-
-    $$
-    g_{syn} = u \cdot x
-    $$
-
-    where:
-    - $u$ represents the utilization of synaptic efficacy (facilitation variable)
-    - $x$ represents the available synaptic resources (depression variable)
-    - $\tau_f$ is the facilitation time constant
-    - $\tau_d$ is the depression time constant
-    - $U$ is the baseline utilization parameter
-    - $\delta(t - t_{spike})$ is the Dirac delta function representing presynaptic spikes
-    - $g_{syn}$ is the effective synaptic conductance
-
-    Parameters
-    ----------
-    in_size : Size
-        Size of the input.
-    name : str, optional
-        Name of the synapse instance.
-    U : ArrayLike, default=0.15
-        Baseline utilization parameter (fraction of resources used per action potential).
-    tau_f : ArrayLike, default=1500.*u.ms
-        Time constant of short-term facilitation in milliseconds.
-    tau_d : ArrayLike, default=200.*u.ms
-        Time constant of short-term depression (recovery of synaptic resources) in milliseconds.
-
-    Attributes
-    ----------
-    u : HiddenState
-        Utilization of synaptic efficacy (facilitation variable).
-    x : HiddenState
-        Available synaptic resources (depression variable).
-
-    Notes
-    -----
-    - Larger values of tau_f produce stronger facilitation effects.
-    - Larger values of tau_d lead to slower recovery from depression.
-    - The parameter U controls the initial release probability.
-    - The effective synaptic strength is the product of u and x.
-
-    References
-    ----------
-    .. [1] Tsodyks, M. V., & Markram, H. (1997). The neural code between neocortical
-           pyramidal neurons depends on neurotransmitter release probability.
-           Proceedings of the National Academy of Sciences, 94(2), 719-723.
-    .. [2] Tsodyks, M., Pawelzik, K., & Markram, H. (1998). Neural networks with dynamic
-           synapses. Neural computation, 10(4), 821-835.
-    """
-    __module__ = 'brainstate.nn'
-
-    def __init__(
-        self,
-        in_size: Size,
-        name: Optional[str] = None,
-        U: ArrayLike = 0.15,
-        tau_f: ArrayLike = 1500. * u.ms,
-        tau_d: ArrayLike = 200. * u.ms,
-    ):
-        super().__init__(name=name, in_size=in_size)
-
-        # parameters
-        self.tau_f = init.param(tau_f, self.varshape)
-        self.tau_d = init.param(tau_d, self.varshape)
-        self.U = init.param(U, self.varshape)
-
-    def init_state(self, batch_size: int = None, **kwargs):
-        self.x = HiddenState(init.param(init.Constant(1.), self.varshape, batch_size))
-        self.u = HiddenState(init.param(init.Constant(self.U), self.varshape, batch_size))
-
-    def reset_state(self, batch_size: int = None, **kwargs):
-        self.x.value = init.param(init.Constant(1.), self.varshape, batch_size)
-        self.u.value = init.param(init.Constant(self.U), self.varshape, batch_size)
-
-    def update(self, pre_spike):
-        u = exp_euler_step(lambda u: - u / self.tau_f, self.u.value)
-        x = exp_euler_step(lambda x: (1 - x) / self.tau_d, self.x.value)
-
-        # --- original code:
-        #   if pre_spike.dtype == jax.numpy.bool_:
-        #     u = bm.where(pre_spike, u + self.U * (1 - self.u), u)
-        #     x = bm.where(pre_spike, x - u * self.x, x)
-        #   else:
-        #     u = pre_spike * (u + self.U * (1 - self.u)) + (1 - pre_spike) * u
-        #     x = pre_spike * (x - u * self.x) + (1 - pre_spike) * x
+        return self.update_return()
 
-        # --- simplified code:
-        u = u + pre_spike * self.U * (1 - self.u.value)
-        x = x - pre_spike * u * self.x.value
-
-        self.u.value = u
-        self.x.value = x
-        return u * x * pre_spike
-
-
-class STD(Synapse):
-    r"""
-    Synapse with short-term depression.
-
-    This class implements a synapse model with short-term depression (STD), which captures
-    activity-dependent reduction in synaptic efficacy, typically caused by depletion of
-    neurotransmitter vesicles following repeated stimulation.
-
-    The model is characterized by the following equation:
-
-    $$
-    \frac{dx}{dt} = \frac{1 - x}{\tau} - U \cdot x \cdot \delta(t - t_{spike})
-    $$
-
-    $$
-    g_{syn} = x
-    $$
-
-    where:
-    - $x$ represents the available synaptic resources (depression variable)
-    - $\tau$ is the depression recovery time constant
-    - $U$ is the utilization parameter (fraction of resources depleted per spike)
-    - $\delta(t - t_{spike})$ is the Dirac delta function representing presynaptic spikes
-    - $g_{syn}$ is the effective synaptic conductance
-
-    Parameters
-    ----------
-    in_size : Size
-        Size of the input.
-    name : str, optional
-        Name of the synapse instance.
-    tau : ArrayLike, default=200.*u.ms
-        Time constant governing recovery of synaptic resources in milliseconds.
-    U : ArrayLike, default=0.07
-        Utilization parameter (fraction of resources used per action potential).
-
-    Attributes
-    ----------
-    x : HiddenState
-        Available synaptic resources (depression variable).
-
-    Notes
-    -----
-    - Larger values of tau lead to slower recovery from depression.
-    - Larger values of U cause stronger depression with each spike.
-    - This model is a simplified version of the STP model that only includes depression.
-
-    References
-    ----------
-    .. [1] Abbott, L. F., Varela, J. A., Sen, K., & Nelson, S. B. (1997). Synaptic
-           depression and cortical gain control. Science, 275(5297), 220-224.
-    .. [2] Tsodyks, M. V., & Markram, H. (1997). The neural code between neocortical
-           pyramidal neurons depends on neurotransmitter release probability.
-           Proceedings of the National Academy of Sciences, 94(2), 719-723.
-    """
-    __module__ = 'brainstate.nn'
-
-    def __init__(
-        self,
-        in_size: Size,
-        name: Optional[str] = None,
-        # synapse parameters
-        tau: ArrayLike = 200. * u.ms,
-        U: ArrayLike = 0.07,
-    ):
-        super().__init__(name=name, in_size=in_size)
-
-        # parameters
-        self.tau = init.param(tau, self.varshape)
-        self.U = init.param(U, self.varshape)
-
-    def init_state(self, batch_size: int = None, **kwargs):
-        self.x = HiddenState(init.param(init.Constant(1.), self.varshape, batch_size))
-
-    def reset_state(self, batch_size: int = None, **kwargs):
-        self.x.value = init.param(init.Constant(1.), self.varshape, batch_size)
-
-    def update(self, pre_spike):
-        x = exp_euler_step(lambda x: (1 - x) / self.tau, self.x.value)
-
-        # --- original code:
-        # self.x.value = bm.where(pre_spike, x - self.U * self.x, x)
-
-        # --- simplified code:
-        self.x.value = x - pre_spike * self.U * self.x.value
-
-        return self.x.value * pre_spike
+    def update_return(self) -> PyTree:
+        return self.g.value
 
 
 class AMPA(Synapse):
@@ -614,6 +421,10 @@ class AMPA(Synapse):
         self.spike_arrival_time.value = u.math.where(pre_spike, t, self.spike_arrival_time.value)
         TT = ((t - self.spike_arrival_time.value) < self.T_duration) * self.T
         self.g.value = exp_euler_step(self.dg, self.g.value, t, TT)
+        return self.update_return()
+
+    def update_return(self) -> PyTree:
+        """Return the synaptic conductance value."""
         return self.g.value