diff --git a/brainpy/math/math_sparse_surrogate_fixes_test.py b/brainpy/math/math_sparse_surrogate_fixes_test.py index e35fd15b0..f943f8d5a 100644 --- a/brainpy/math/math_sparse_surrogate_fixes_test.py +++ b/brainpy/math/math_sparse_surrogate_fixes_test.py @@ -16,11 +16,11 @@ (``csr_to_dense`` wraps ``brainevent.CSR(...).todense()`` correctly). * ``brainpy/math/jitconn/matvec.py`` — M-13 (``mv_prob_*`` / ``event_mv_prob_*`` are reproducible when an explicit ``seed`` is threaded). -* ``brainpy/math/surrogate/_one_input.py`` and - ``brainpy/math/surrogate/_one_input_new.py`` — H-20..H-24: for every surrogate - the ``surrogate_grad`` matches ``jax.grad(surrogate_fun)``; ``GaussianGrad`` - widens with ``sigma`` (H-20); ``PiecewiseQuadratic`` grad matches its forward - derivative (H-21); ``QPseudoSpike`` uses ``alpha-1`` (H-22); ``Arctan`` +* ``brainpy.math.surrogate`` (the local package was removed) — now an alias of + ``braintools.surrogate`` (>=0.2.0). H-20..H-24: for every reused surrogate the + ``surrogate_grad`` matches ``jax.grad(surrogate_fun)``; ``GaussianGrad`` widens + with ``sigma`` (H-20); ``PiecewiseQuadratic`` grad matches its forward + derivative (H-21); ``QPseudoSpike`` grad at 0 is 1 (H-22); ``Arctan`` ``surrogate_fun`` does not raise (H-23); ``ERF`` ``surrogate_fun`` is increasing (H-24). * ``brainpy/math/delayvars.py`` — C-09 (``TimeDelay`` ring-buffer read @@ -40,8 +40,7 @@ import pytest import brainpy.math as bm -import brainpy.math.surrogate._one_input as old_surr -import brainpy.math.surrogate._one_input_new as new_surr +import braintools.surrogate as _bt_surrogate # --------------------------------------------------------------------------- @@ -248,42 +247,41 @@ def test_event_csrmv_matches_masked_dense(): # =========================================================================== -# H-20..H-24 — surrogate gradients consistent with their forward functions +# Surrogate gradients now reuse ``braintools.surrogate`` (>=0.2.0). +# +# The local ``brainpy/math/surrogate`` package was removed; ``bm.surrogate`` is +# an alias of ``braintools.surrogate``. These tests lock in (a) the re-export +# wiring and (b) the audit correctness properties H-20..H-24 against the reused, +# fixed braintools implementation, so a braintools regression / accidental +# downgrade is caught here. # =========================================================================== -# The two modules expose slightly different APIs: -# _one_input.py : surrogate_grad(self, dz, x) (dz = upstream gradient) -# _one_input_new.py : surrogate_grad(self, x) -# These helpers normalise that so the same assertions run against both. - -def _grad_old(inst, x): - return inst.surrogate_grad(1.0, x) - - -def _grad_new(inst, x): - return inst.surrogate_grad(x) - +def test_surrogate_is_braintools_reexport(): + # bm.surrogate is the braintools module itself (no local duplicate left). + assert bm.surrogate is _bt_surrogate + # neuron defaults reference these names; they must resolve and be callable. + assert callable(bm.surrogate.InvSquareGrad) + assert callable(bm.surrogate.relu_grad) + assert callable(bm.surrogate.InvSquareGrad()) -_MODULES = [ - ("_one_input", old_surr, _grad_old), - ("_one_input_new", new_surr, _grad_new), -] -# Surrogates that expose BOTH surrogate_fun and surrogate_grad and are -# differentiable away from kinks — used for the grad-vs-autograd check. +# Surrogates exposing BOTH surrogate_fun and surrogate_grad, differentiable away +# from kinks -- used for the grad-vs-autograd self-consistency check. _HAS_FUN = ["PiecewiseQuadratic", "QPseudoSpike", "Arctan", "ERF"] +# Class / functional names re-exported from braintools. +_ONE_INPUT_CLASSES = [n for n in _bt_surrogate.__all__ + if n[0].isupper() and n != "Surrogate"] +_FUNC_NAMES = [n for n in _bt_surrogate.__all__ if n[0].islower()] + -@pytest.mark.parametrize("modname,mod,getgrad", _MODULES, - ids=[m[0] for m in _MODULES]) @pytest.mark.parametrize("clsname", _HAS_FUN) -def test_surrogate_grad_matches_autograd(modname, mod, getgrad, clsname): - """H-21/H-22/H-23/H-24: surrogate_grad == d/dx surrogate_fun.""" - cls = getattr(mod, clsname) - inst = cls() +def test_surrogate_grad_matches_autograd(clsname): + """H-21..H-24: surrogate_grad == d/dx surrogate_fun on the reused impls.""" + inst = getattr(bm.surrogate, clsname)() # Avoid the exact kinks (|x| = 1/alpha) where the piecewise derivative jumps. xs = jnp.asarray([-0.85, -0.4, -0.1, 0.1, 0.4, 0.85], dtype=jnp.float32) - analytic = np.asarray(getgrad(inst, xs)) + analytic = np.asarray(inst.surrogate_grad(xs)) def fun_scalar(v): return jnp.squeeze(inst.surrogate_fun(jnp.reshape(v, (1,)))) @@ -292,184 +290,70 @@ def fun_scalar(v): np.testing.assert_allclose(analytic, autograd, rtol=2e-3, atol=2e-4) -@pytest.mark.parametrize("modname,mod,getgrad", _MODULES, - ids=[m[0] for m in _MODULES]) -@pytest.mark.parametrize("clsname", _HAS_FUN) -def test_surrogate_fun_monotone_increasing_on_unit_interval(modname, mod, getgrad, clsname): - """Each surrogate forward (origin) function is non-decreasing on [0, 1].""" - cls = getattr(mod, clsname) - inst = cls() - fx = np.asarray(inst.surrogate_fun(jnp.linspace(0.0, 1.0, 21))) - assert np.all(np.diff(fx) >= -1e-6) - - -@pytest.mark.parametrize("modname,mod,getgrad", _MODULES, - ids=[m[0] for m in _MODULES]) -def test_arctan_surrogate_fun_does_not_raise(modname, mod, getgrad): - """H-23: Arctan.surrogate_fun previously called jnp.arctan2 with one arg.""" - inst = mod.Arctan() - out = np.asarray(inst.surrogate_fun(jnp.asarray([-0.5, -0.1, 0.0, 0.1, 0.5]))) +def test_arctan_surrogate_fun_does_not_raise(): + """H-23: Arctan.surrogate_fun is finite, increasing, and centred at 0.5.""" + out = np.asarray(bm.surrogate.Arctan().surrogate_fun( + jnp.asarray([-0.5, -0.1, 0.0, 0.1, 0.5]))) assert np.all(np.isfinite(out)) - # arctan forward crosses 0.5 at x = 0 - assert np.isclose(out[2], 0.5, atol=1e-6) + assert np.isclose(out[2], 0.5, atol=1e-6) # arctan forward crosses 0.5 at x=0 assert np.all(np.diff(out) > 0) -@pytest.mark.parametrize("modname,mod,getgrad", _MODULES, - ids=[m[0] for m in _MODULES]) -def test_erf_surrogate_fun_is_increasing(modname, mod, getgrad): - """H-24: ERF.surrogate_fun must be increasing (was decreasing before).""" - inst = mod.ERF() - out = np.asarray(inst.surrogate_fun(jnp.linspace(-0.5, 0.5, 11))) +def test_erf_surrogate_fun_is_increasing(): + """H-24: ERF.surrogate_fun must be increasing and centred at 0.5.""" + out = np.asarray(bm.surrogate.ERF().surrogate_fun(jnp.linspace(-0.5, 0.5, 11))) assert np.all(np.diff(out) > 0) - assert np.isclose(out[5], 0.5, atol=1e-6) # centred at x = 0 + assert np.isclose(out[5], 0.5, atol=1e-6) -@pytest.mark.parametrize("modname,mod,getgrad", _MODULES, - ids=[m[0] for m in _MODULES]) -def test_gaussian_grad_bump_widens_with_sigma(modname, mod, getgrad): - """H-20: GaussianGrad — at x=1 the gradient must INCREASE with sigma - (a wider bump), proving the sigma is no longer inverted by the - operator-precedence bug ``exp(-(x**2)/2*sigma**2)``.""" - g_narrow = float(np.asarray(getgrad(mod.GaussianGrad(sigma=0.5), jnp.asarray(1.0)))) - g_wide = float(np.asarray(getgrad(mod.GaussianGrad(sigma=2.0), jnp.asarray(1.0)))) +def test_gaussian_grad_bump_widens_with_sigma(): + """H-20: GaussianGrad gradient at x=1 INCREASES with sigma (a wider bump), + proving sigma is not inverted by the operator-precedence bug.""" + g_narrow = float(np.asarray( + bm.surrogate.GaussianGrad(sigma=0.5).surrogate_grad(jnp.asarray(1.0)))) + g_wide = float(np.asarray( + bm.surrogate.GaussianGrad(sigma=2.0).surrogate_grad(jnp.asarray(1.0)))) assert g_wide > g_narrow - # Sanity on the intended magnitude (audit: grad@±1 ≈ 0.088 for sigma=2). - assert g_wide == pytest.approx(0.088, abs=2e-2) - - -@pytest.mark.parametrize("modname,mod,getgrad", _MODULES, - ids=[m[0] for m in _MODULES]) -def test_piecewise_quadratic_grad_formula(modname, mod, getgrad): - """H-21: grad == -alpha**2 |x| + alpha inside the support, 0 outside.""" - inst = mod.PiecewiseQuadratic(alpha=1.0) - g_in = float(np.asarray(getgrad(inst, jnp.asarray(0.5)))) - assert g_in == pytest.approx(-1.0 * 0.5 + 1.0) # = 0.5 - g_out = float(np.asarray(getgrad(inst, jnp.asarray(5.0)))) - assert g_out == pytest.approx(0.0) - - -@pytest.mark.parametrize("modname,mod,getgrad", _MODULES, - ids=[m[0] for m in _MODULES]) -def test_qpseudospike_grad_uses_alpha_minus_one(modname, mod, getgrad): - """H-22: grad denominator uses (alpha-1); grad at 0 == 1.""" - inst = mod.QPseudoSpike(alpha=2.0) - g0 = float(np.asarray(getgrad(inst, jnp.asarray(0.0)))) - assert g0 == pytest.approx(1.0, abs=1e-6) - + assert g_wide == pytest.approx(0.088, abs=2e-2) # audit: grad@1 ~ 0.088 (sigma=2) -# =========================================================================== -# Surrogate coverage — every class's __call__ + surrogate_grad in both modules -# =========================================================================== - -def _new_surrogate_classes(): - return [getattr(new_surr, n) for n in new_surr.__all__ - if n[0].isupper() and n != "Surrogate"] +def test_piecewise_quadratic_grad_formula(): + """H-21: grad == alpha - alpha**2 |x| inside the support, 0 outside.""" + inst = bm.surrogate.PiecewiseQuadratic(alpha=1.0) + assert float(np.asarray(inst.surrogate_grad(jnp.asarray(0.5)))) == pytest.approx(0.5) + assert float(np.asarray(inst.surrogate_grad(jnp.asarray(5.0)))) == pytest.approx(0.0) -def _old_surrogate_classes(): - out = [] - for n in dir(old_surr): - obj = getattr(old_surr, n) - if (isinstance(obj, type) and issubclass(obj, old_surr.Surrogate) - and n not in ("Surrogate", "_OneInpSurrogate")): - out.append(obj) - return out - -@pytest.mark.parametrize("cls", _new_surrogate_classes(), - ids=lambda c: c.__name__) -def test_new_surrogate_call_and_grad_run(cls): - inst = cls() - x = jnp.linspace(-1.5, 1.5, 9) - y = inst(x) # __call__ -> heaviside forward - assert np.asarray(y).shape == (9,) - # forward is a {0,1} spike indicator - assert set(np.unique(np.asarray(y)).tolist()).issubset({0.0, 1.0}) - g = np.asarray(inst.surrogate_grad(x)) # surrogate_grad(x) - assert g.shape == (9,) and np.all(np.isfinite(g)) - # grad flows through __call__ - flow = jax.grad(lambda v: jnp.sum(inst(v)))(x) - assert np.all(np.isfinite(np.asarray(flow))) +def test_qpseudospike_grad_at_zero_is_one(): + """H-22: q-PseudoSpike gradient at 0 == 1.""" + g0 = float(np.asarray( + bm.surrogate.QPseudoSpike(alpha=2.0).surrogate_grad(jnp.asarray(0.0)))) + assert g0 == pytest.approx(1.0, abs=1e-6) -@pytest.mark.parametrize("cls", _old_surrogate_classes(), - ids=lambda c: c.__name__) -def test_old_surrogate_call_and_grad_run(cls): - inst = cls() +@pytest.mark.parametrize("clsname", _ONE_INPUT_CLASSES, ids=lambda n: n) +def test_surrogate_call_and_grad_run(clsname): + """Every re-exported class: __call__ is a {0,1} spike + grad flows finite.""" + inst = getattr(bm.surrogate, clsname)() x = jnp.linspace(-1.5, 1.5, 9) - y = inst(x) # custom-gradient forward + y = inst(x) # __call__ -> heaviside forward assert np.asarray(y).shape == (9,) assert set(np.unique(np.asarray(y)).tolist()).issubset({0.0, 1.0}) - g = np.asarray(inst.surrogate_grad(1.0, x)) # surrogate_grad(dz, x) + g = np.asarray(inst.surrogate_grad(x)) assert g.shape == (9,) and np.all(np.isfinite(g)) flow = jax.grad(lambda v: jnp.sum(inst(v)))(x) assert np.all(np.isfinite(np.asarray(flow))) -def test_new_surrogate_repr_and_functional_aliases(): - # Exercise functional (lowercase) entry points + __repr__ for coverage. - x = jnp.linspace(-1.0, 1.0, 5) - assert "Arctan" in repr(new_surr.Arctan()) - for fn in (new_surr.sigmoid, new_surr.arctan, new_surr.erf, - new_surr.gaussian_grad, new_surr.relu_grad): - assert np.asarray(fn(x)).shape == (5,) - - -def test_old_surrogate_repr_and_functional_aliases(): - x = jnp.linspace(-1.0, 1.0, 5) - assert "GaussianGrad" in repr(old_surr.GaussianGrad()) - for fn in (old_surr.sigmoid, old_surr.arctan, old_surr.erf, - old_surr.gaussian_grad, old_surr.q_pseudo_spike): - assert np.asarray(fn(x)).shape == (5,) - - -# Lowercase functional aliases present in BOTH modules. -_FUNC_NAMES = [n for n in new_surr.__all__ if n[0].islower()] - - -@pytest.mark.parametrize("fname", _FUNC_NAMES) -def test_old_functional_alias_forward_and_origin(fname): - """Exercise every ``_one_input`` functional alias (heaviside forward and, - where supported, the ``origin=True`` smooth forward).""" - import inspect - fn = getattr(old_surr, fname) - x = jnp.linspace(-1.2, 1.2, 7) - y = np.asarray(fn(x)) - assert y.shape == (7,) and np.all(np.isfinite(y)) - if "origin" in inspect.signature(fn).parameters: - yo = np.asarray(fn(x, origin=True)) # exercises surrogate_fun - assert yo.shape == (7,) and np.all(np.isfinite(yo)) - - @pytest.mark.parametrize("fname", _FUNC_NAMES) -def test_new_functional_alias_forward(fname): - """Exercise every ``_one_input_new`` functional alias (heaviside forward).""" - fn = getattr(new_surr, fname) +def test_functional_alias_forward(fname): + """Every re-exported functional alias returns a finite heaviside forward.""" + fn = getattr(bm.surrogate, fname) x = jnp.linspace(-1.2, 1.2, 7) y = np.asarray(fn(x)) assert y.shape == (7,) and np.all(np.isfinite(y)) -def _new_classes_with_surrogate_fun(): - out = [] - for n in new_surr.__all__: - if not (n[0].isupper() and n != "Surrogate"): - continue - c = getattr(new_surr, n) - if c.surrogate_fun is not new_surr.Surrogate.surrogate_fun: - out.append(c) - return out - - -@pytest.mark.parametrize("cls", _new_classes_with_surrogate_fun(), - ids=lambda c: c.__name__) -def test_new_surrogate_fun_runs(cls): - """Cover the ``surrogate_fun`` body of every new-module class that has one.""" - out = np.asarray(cls().surrogate_fun(jnp.linspace(-1.2, 1.2, 9))) - assert out.shape == (9,) and np.all(np.isfinite(out)) - - # =========================================================================== # C-09 — TimeDelay ring-buffer read applies the modulo # =========================================================================== diff --git a/brainpy/math/surrogate/__init__.py b/brainpy/math/surrogate/__init__.py deleted file mode 100644 index 00f6a61b7..000000000 --- a/brainpy/math/surrogate/__init__.py +++ /dev/null @@ -1,17 +0,0 @@ -# -*- coding: utf-8 -*- -# Copyright 2025 BrainX 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 ._one_input_new import * -from ._two_inputs import * diff --git a/brainpy/math/surrogate/_one_input.py b/brainpy/math/surrogate/_one_input.py deleted file mode 100644 index 6f6ac5162..000000000 --- a/brainpy/math/surrogate/_one_input.py +++ /dev/null @@ -1,1749 +0,0 @@ -# -*- coding: utf-8 -*- -# Copyright 2025 BrainX 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. -# ============================================================================== -import functools -from typing import Union - -import jax -import jax.numpy as jnp -import jax.scipy as sci - -from brainpy.math.interoperability import as_jax -from brainpy.math.ndarray import Array - -__all__ = [ - 'sigmoid', - 'piecewise_quadratic', - 'piecewise_exp', - 'soft_sign', - 'arctan', - 'nonzero_sign_log', - 'erf', - 'piecewise_leaky_relu', - 'squarewave_fourier_series', - 's2nn', - 'q_pseudo_spike', - 'leaky_relu', - 'log_tailed_relu', - 'relu_grad', - 'gaussian_grad', - 'inv_square_grad', - 'multi_gaussian_grad', - 'slayer_grad', -] - - -class Surrogate(object): - """The base surrograte gradient function.""" - - def __call__(self, *args, **kwargs): - raise NotImplementedError - - def __repr__(self): - return f'{self.__class__.__name__}()' - - -class _OneInpSurrogate(Surrogate): - def __init__(self, forward_use_surrogate=False): - self.forward_use_surrogate = forward_use_surrogate - self._true_call_ = jax.custom_gradient(self.call) - - def __call__(self, x: jax.Array): - return self._true_call_(as_jax(x)) - - def call(self, x): - """Call the function for surrogate gradient propagation.""" - y = self.surrogate_fun(x) if self.forward_use_surrogate else self.true_fun(x) - return y, functools.partial(self.surrogate_grad, x=x) - - def true_fun(self, x): - """The original true function.""" - return jnp.asarray(x >= 0, dtype=x.dtype) - - def surrogate_fun(self, x): - """The surrogate function.""" - raise NotImplementedError - - def surrogate_grad(self, dz, x): - """The gradient for the surrogate function.""" - raise NotImplementedError - - def __repr__(self): - return f'{self.__class__.__name__}(forward_use_surrogate={self.forward_use_surrogate})' - - -class Sigmoid(_OneInpSurrogate): - """Spike function with the sigmoid-shaped surrogate gradient. - - See Also:: - - sigmoid - - """ - - def __init__(self, alpha: float = 4., forward_use_surrogate=False): - super().__init__(forward_use_surrogate) - self.alpha = alpha - - def surrogate_fun(self, x): - return sci.special.expit(x) - - def surrogate_grad(self, dz, x): - sgax = sci.special.expit(as_jax(x) * self.alpha) - dx = as_jax(dz) * (1. - sgax) * sgax * self.alpha - return dx - - def __repr__(self): - return f'{self.__class__.__name__}(alpha={self.alpha})' - - -def sigmoid( - x: Union[jax.Array, Array], - alpha: float = 4., - origin: bool = False, -): - r"""Spike function with the sigmoid-shaped surrogate gradient. - - If `origin=False`, return the forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - If `origin=True`, computes the original function: - - .. math:: - - g(x) = \mathrm{sigmoid}(\alpha x) = \frac{1}{1+e^{-\alpha x}} - - Backward function: - - .. math:: - - g'(x) = \alpha * (1 - \mathrm{sigmoid} (\alpha x)) \mathrm{sigmoid} (\alpha x) - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> bp.visualize.get_figure(1, 1, 4, 6) - >>> xs = bm.linspace(-2, 2, 1000) - >>> for alpha in [1., 2., 4.]: - >>> grads = bm.vector_grad(bm.surrogate.sigmoid)(xs, alpha) - >>> plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha)) - >>> plt.legend() - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - alpha: float - Parameter to control smoothness of gradient - origin: bool - Whether to compute the original function as the feedfoward output. - - Returns:: - - out: jax.Array - The spiking state. - """ - return Sigmoid(alpha=alpha, forward_use_surrogate=origin)(x) - - -class PiecewiseQuadratic(_OneInpSurrogate): - """Judge spiking state with a piecewise quadratic function. - - See Also:: - - piecewise_quadratic - - """ - - def __init__(self, alpha: float = 1., forward_use_surrogate=False): - super().__init__(forward_use_surrogate) - self.alpha = alpha - - def surrogate_fun(self, x): - x = as_jax(x) - z = jnp.where(x < -1 / self.alpha, - 0., - jnp.where(x > 1 / self.alpha, - 1., - (-self.alpha * jnp.abs(x) / 2 + 1) * self.alpha * x + 0.5)) - return z - - def surrogate_grad(self, dz, x): - x = as_jax(x) - dx = jnp.where(jnp.abs(x) > 1 / self.alpha, 0., dz * (-self.alpha ** 2 * jnp.abs(x) + self.alpha)) - return dx - - def __repr__(self): - return f'{self.__class__.__name__}(alpha={self.alpha})' - - -def piecewise_quadratic( - x: Union[jax.Array, Array], - alpha: float = 1., - origin: bool = False -): - r"""Judge spiking state with a piecewise quadratic function [1]_ [2]_ [3]_ [4]_ [5]_. - - If `origin=False`, computes the forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - If `origin=True`, computes the original function: - - .. math:: - - g(x) = - \begin{cases} - 0, & x < -\frac{1}{\alpha} \\ - -\frac{1}{2}\alpha^2|x|x + \alpha x + \frac{1}{2}, & |x| \leq \frac{1}{\alpha} \\ - 1, & x > \frac{1}{\alpha} \\ - \end{cases} - - Backward function: - - .. math:: - - g'(x) = - \begin{cases} - 0, & |x| > \frac{1}{\alpha} \\ - -\alpha^2|x|+\alpha, & |x| \leq \frac{1}{\alpha} - \end{cases} - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> bp.visualize.get_figure(1, 1, 4, 6) - >>> xs = bm.linspace(-3, 3, 1000) - >>> for alpha in [0.5, 1., 2., 4.]: - >>> grads = bm.vector_grad(bm.surrogate.piecewise_quadratic)(xs, alpha) - >>> plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha)) - >>> plt.legend() - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - alpha: float - Parameter to control smoothness of gradient - origin: bool - Whether to compute the original function as the feedfoward output. - - Returns:: - - out: jax.Array - The spiking state. - - References:: - - .. [1] Esser S K, Merolla P A, Arthur J V, et al. Convolutional networks for fast, energy-efficient neuromorphic computing[J]. Proceedings of the national academy of sciences, 2016, 113(41): 11441-11446. - .. [2] Wu Y, Deng L, Li G, et al. Spatio-temporal backpropagation for training high-performance spiking neural networks[J]. Frontiers in neuroscience, 2018, 12: 331. - .. [3] Bellec G, Salaj D, Subramoney A, et al. Long short-term memory and learning-to-learn in networks of spiking neurons[C]//Proceedings of the 32nd International Conference on Neural Information Processing Systems. 2018: 795-805. - .. [4] Neftci E O, Mostafa H, Zenke F. Surrogate gradient learning in spiking neural networks: Bringing the power of gradient-based optimization to spiking neural networks[J]. IEEE Signal Processing Magazine, 2019, 36(6): 51-63. - .. [5] Panda P, Aketi S A, Roy K. Toward scalable, efficient, and accurate deep spiking neural networks with backward residual connections, stochastic softmax, and hybridization[J]. Frontiers in Neuroscience, 2020, 14. - """ - return PiecewiseQuadratic(alpha=alpha, forward_use_surrogate=origin)(x) - - -class PiecewiseExp(_OneInpSurrogate): - """Judge spiking state with a piecewise exponential function. - - See Also:: - - piecewise_exp - """ - - def __init__(self, alpha: float = 1., forward_use_surrogate=False): - super().__init__(forward_use_surrogate) - self.alpha = alpha - - def surrogate_grad(self, dz, x): - x = as_jax(x) - dx = (self.alpha / 2) * jnp.exp(-self.alpha * jnp.abs(x)) - return dx * as_jax(dz) - - def surrogate_fun(self, x): - x = as_jax(x) - return jnp.where(x < 0, jnp.exp(self.alpha * x) / 2, 1 - jnp.exp(-self.alpha * x) / 2) - - def __repr__(self): - return f'{self.__class__.__name__}(alpha={self.alpha})' - - -def piecewise_exp( - x: Union[jax.Array, Array], - alpha: float = 1., - origin: bool = False -): - r"""Judge spiking state with a piecewise exponential function [1]_. - - If `origin=False`, computes the forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - If `origin=True`, computes the original function: - - .. math:: - - g(x) = \begin{cases} - \frac{1}{2}e^{\alpha x}, & x < 0 \\ - 1 - \frac{1}{2}e^{-\alpha x}, & x \geq 0 - \end{cases} - - Backward function: - - .. math:: - - g'(x) = \frac{\alpha}{2}e^{-\alpha |x|} - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> bp.visualize.get_figure(1, 1, 4, 6) - >>> xs = bm.linspace(-3, 3, 1000) - >>> for alpha in [0.5, 1., 2., 4.]: - >>> grads = bm.vector_grad(bm.surrogate.piecewise_exp)(xs, alpha) - >>> plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha)) - >>> plt.legend() - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - alpha: float - Parameter to control smoothness of gradient - origin: bool - Whether to compute the original function as the feedfoward output. - - Returns:: - - out: jax.Array - The spiking state. - - References:: - - .. [1] Neftci E O, Mostafa H, Zenke F. Surrogate gradient learning in spiking neural networks: Bringing the power of gradient-based optimization to spiking neural networks[J]. IEEE Signal Processing Magazine, 2019, 36(6): 51-63. - """ - return PiecewiseExp(alpha=alpha, forward_use_surrogate=origin)(x) - - -class SoftSign(_OneInpSurrogate): - """Judge spiking state with a soft sign function. - - See Also:: - - soft_sign - """ - - def __init__(self, alpha=1., forward_use_surrogate=False): - super().__init__(forward_use_surrogate=forward_use_surrogate) - self.alpha = alpha - - def surrogate_grad(self, dz, x): - x = as_jax(x) - dx = self.alpha * 0.5 / (1 + jnp.abs(self.alpha * x)) ** 2 - return dx * as_jax(dz) - - def surrogate_fun(self, x): - x = as_jax(x) - return x / (2 / self.alpha + 2 * jnp.abs(x)) + 0.5 - - def __repr__(self): - return f'{self.__class__.__name__}(alpha={self.alpha})' - - -def soft_sign( - x: Union[jax.Array, Array], - alpha: float = 1., - origin: bool = False -): - r"""Judge spiking state with a soft sign function. - - If `origin=False`, computes the forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - If `origin=True`, computes the original function: - - .. math:: - - g(x) = \frac{1}{2} (\frac{\alpha x}{1 + |\alpha x|} + 1) - = \frac{1}{2} (\frac{x}{\frac{1}{\alpha} + |x|} + 1) - - Backward function: - - .. math:: - - g'(x) = \frac{\alpha}{2(1 + |\alpha x|)^{2}} = \frac{1}{2\alpha(\frac{1}{\alpha} + |x|)^{2}} - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> bp.visualize.get_figure(1, 1, 4, 6) - >>> xs = bm.linspace(-3, 3, 1000) - >>> for alpha in [0.5, 1., 2., 4.]: - >>> grads = bm.vector_grad(bm.surrogate.soft_sign)(xs, alpha) - >>> plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha)) - >>> plt.legend() - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - alpha: float - Parameter to control smoothness of gradient - origin: bool - Whether to compute the original function as the feedfoward output. - - Returns:: - - out: jax.Array - The spiking state. - - """ - return SoftSign(alpha=alpha, forward_use_surrogate=origin)(x) - - -class Arctan(_OneInpSurrogate): - """Judge spiking state with an arctan function. - - See Also:: - - arctan - """ - - def __init__(self, alpha=1., forward_use_surrogate=False): - super().__init__(forward_use_surrogate=forward_use_surrogate) - self.alpha = alpha - - def surrogate_grad(self, dz, x): - x = as_jax(x) - dx = self.alpha * 0.5 / (1 + (jnp.pi / 2 * self.alpha * x) ** 2) - return dx * as_jax(dz) - - def surrogate_fun(self, x): - x = as_jax(x) - return jnp.arctan(jnp.pi / 2 * self.alpha * x) / jnp.pi + 0.5 - - def __repr__(self): - return f'{self.__class__.__name__}(alpha={self.alpha})' - - -def arctan( - x: Union[jax.Array, Array], - alpha: float = 1., - origin: bool = False -): - r"""Judge spiking state with an arctan function. - - If `origin=False`, computes the forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - If `origin=True`, computes the original function: - - .. math:: - - g(x) = \frac{1}{\pi} \arctan(\frac{\pi}{2}\alpha x) + \frac{1}{2} - - Backward function: - - .. math:: - - g'(x) = \frac{\alpha}{2(1 + (\frac{\pi}{2}\alpha x)^2)} - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> bp.visualize.get_figure(1, 1, 4, 6) - >>> xs = bm.linspace(-3, 3, 1000) - >>> for alpha in [0.5, 1., 2., 4.]: - >>> grads = bm.vector_grad(bm.surrogate.arctan)(xs, alpha) - >>> plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha)) - >>> plt.legend() - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - alpha: float - Parameter to control smoothness of gradient - origin: bool - Whether to compute the original function as the feedfoward output. - - Returns:: - - out: jax.Array - The spiking state. - - """ - return Arctan(alpha=alpha, forward_use_surrogate=origin)(x) - - -class NonzeroSignLog(_OneInpSurrogate): - """Judge spiking state with a nonzero sign log function. - - See Also:: - - nonzero_sign_log - """ - - def __init__(self, alpha=1., forward_use_surrogate=False): - super().__init__(forward_use_surrogate=forward_use_surrogate) - self.alpha = alpha - - def surrogate_grad(self, dz, x): - x = as_jax(x) - dx = as_jax(dz) / (1 / self.alpha + jnp.abs(x)) - return dx - - def surrogate_fun(self, x): - x = as_jax(x) - return jnp.where(x < 0, -1., 1.) * jnp.log(jnp.abs(self.alpha * x) + 1) - - def __repr__(self): - return f'{self.__class__.__name__}(alpha={self.alpha})' - - -def nonzero_sign_log( - x: Union[jax.Array, Array], - alpha: float = 1., - origin: bool = False -): - r"""Judge spiking state with a nonzero sign log function. - - If `origin=False`, computes the forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - If `origin=True`, computes the original function: - - .. math:: - - g(x) = \mathrm{NonzeroSign}(x) \log (|\alpha x| + 1) - - where - - .. math:: - - \begin{split}\mathrm{NonzeroSign}(x) = - \begin{cases} - 1, & x \geq 0 \\ - -1, & x < 0 \\ - \end{cases}\end{split} - - Backward function: - - .. math:: - - g'(x) = \frac{\alpha}{1 + |\alpha x|} = \frac{1}{\frac{1}{\alpha} + |x|} - - This surrogate function has the advantage of low computation cost during the backward. - - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> bp.visualize.get_figure(1, 1, 4, 6) - >>> xs = bm.linspace(-3, 3, 1000) - >>> for alpha in [0.5, 1., 2., 4.]: - >>> grads = bm.vector_grad(bm.surrogate.nonzero_sign_log)(xs, alpha) - >>> plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha)) - >>> plt.legend() - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - alpha: float - Parameter to control smoothness of gradient - origin: bool - Whether to compute the original function as the feedfoward output. - - Returns:: - - out: jax.Array - The spiking state. - - """ - return NonzeroSignLog(alpha=alpha, forward_use_surrogate=origin)(x) - - -class ERF(_OneInpSurrogate): - """Judge spiking state with an erf function. - - See Also:: - - erf - """ - - def __init__(self, alpha=1., forward_use_surrogate=False): - super().__init__(forward_use_surrogate=forward_use_surrogate) - self.alpha = alpha - - def surrogate_grad(self, dz, x): - x = as_jax(x) - dx = (self.alpha / jnp.sqrt(jnp.pi)) * jnp.exp(-jnp.power(self.alpha, 2) * x * x) - return dx * as_jax(dz) - - def surrogate_fun(self, x): - x = as_jax(x) - return 0.5 * (1. - sci.special.erf(-self.alpha * x)) - - def __repr__(self): - return f'{self.__class__.__name__}(alpha={self.alpha})' - - -def erf( - x: Union[jax.Array, Array], - alpha: float = 1., - origin: bool = False -): - r"""Judge spiking state with an erf function [1]_ [2]_ [3]_. - - If `origin=False`, computes the forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - If `origin=True`, computes the original function: - - .. math:: - - \begin{split} - g(x) &= \frac{1}{2}(1-\text{erf}(-\alpha x)) \\ - &= \frac{1}{2} \text{erfc}(-\alpha x) \\ - &= \frac{1}{\sqrt{\pi}}\int_{-\infty}^{\alpha x}e^{-t^2}dt - \end{split} - - Backward function: - - .. math:: - - g'(x) = \frac{\alpha}{\sqrt{\pi}}e^{-\alpha^2x^2} - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> bp.visualize.get_figure(1, 1, 4, 6) - >>> xs = bm.linspace(-3, 3, 1000) - >>> for alpha in [0.5, 1., 2., 4.]: - >>> grads = bm.vector_grad(bm.surrogate.nonzero_sign_log)(xs, alpha) - >>> plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha)) - >>> plt.legend() - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - alpha: float - Parameter to control smoothness of gradient - origin: bool - Whether to compute the original function as the feedfoward output. - - Returns:: - - out: jax.Array - The spiking state. - - References:: - - .. [1] Esser S K, Appuswamy R, Merolla P, et al. Backpropagation for energy-efficient neuromorphic computing[J]. Advances in neural information processing systems, 2015, 28: 1117-1125. - .. [2] Wu Y, Deng L, Li G, et al. Spatio-temporal backpropagation for training high-performance spiking neural networks[J]. Frontiers in neuroscience, 2018, 12: 331. - .. [3] Yin B, Corradi F, Bohté S M. Effective and efficient computation with multiple-timescale spiking recurrent neural networks[C]//International Conference on Neuromorphic Systems 2020. 2020: 1-8. - - """ - return ERF(alpha=alpha, forward_use_surrogate=origin)(x) - - -class PiecewiseLeakyRelu(_OneInpSurrogate): - """Judge spiking state with a piecewise leaky relu function. - - See Also:: - - piecewise_leaky_relu - """ - - def __init__(self, c=0.01, w=1., forward_use_surrogate=False): - super().__init__(forward_use_surrogate=forward_use_surrogate) - self.c = c - self.w = w - - def surrogate_fun(self, x): - x = as_jax(x) - z = jnp.where(x < -self.w, - self.c * x + self.c * self.w, - jnp.where(x > self.w, - self.c * x - self.c * self.w + 1, - 0.5 * x / self.w + 0.5)) - return z - - def surrogate_grad(self, dz, x): - x = as_jax(x) - dx = jnp.where(jnp.abs(x) > self.w, self.c, 1 / self.w) - return dx * as_jax(dz) - - def __repr__(self): - return f'{self.__class__.__name__}(c={self.c}, w={self.w})' - - -def piecewise_leaky_relu( - x: Union[jax.Array, Array], - c: float = 0.01, - w: float = 1., - origin: bool = False -): - r"""Judge spiking state with a piecewise leaky relu function [1]_ [2]_ [3]_ [4]_ [5]_ [6]_ [7]_ [8]_. - - If `origin=False`, computes the forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - If `origin=True`, computes the original function: - - .. math:: - - \begin{split}g(x) = - \begin{cases} - cx + cw, & x < -w \\ - \frac{1}{2w}x + \frac{1}{2}, & -w \leq x \leq w \\ - cx - cw + 1, & x > w \\ - \end{cases}\end{split} - - Backward function: - - .. math:: - - \begin{split}g'(x) = - \begin{cases} - \frac{1}{w}, & |x| \leq w \\ - c, & |x| > w - \end{cases}\end{split} - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> bp.visualize.get_figure(1, 1, 4, 6) - >>> xs = bm.linspace(-3, 3, 1000) - >>> for c in [0.01, 0.05, 0.1]: - >>> for w in [1., 2.]: - >>> grads1 = bm.vector_grad(bm.surrogate.piecewise_leaky_relu)(xs, c=c, w=w) - >>> plt.plot(bm.as_numpy(xs), bm.as_numpy(grads1), label=f'x={c}, w={w}') - >>> plt.legend() - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - c: float - When :math:`|x| > w` the gradient is `c`. - w: float - When :math:`|x| <= w` the gradient is `1 / w`. - origin: bool - Whether to compute the original function as the feedfoward output. - - Returns:: - - out: jax.Array - The spiking state. - - References:: - - .. [1] Yin S, Venkataramanaiah S K, Chen G K, et al. Algorithm and hardware design of discrete-time spiking neural networks based on back propagation with binary activations[C]//2017 IEEE Biomedical Circuits and Systems Conference (BioCAS). IEEE, 2017: 1-5. - .. [2] Wu Y, Deng L, Li G, et al. Spatio-temporal backpropagation for training high-performance spiking neural networks[J]. Frontiers in neuroscience, 2018, 12: 331. - .. [3] Huh D, Sejnowski T J. Gradient descent for spiking neural networks[C]//Proceedings of the 32nd International Conference on Neural Information Processing Systems. 2018: 1440-1450. - .. [4] Wu Y, Deng L, Li G, et al. Direct training for spiking neural networks: Faster, larger, better[C]//Proceedings of the AAAI Conference on Artificial Intelligence. 2019, 33(01): 1311-1318. - .. [5] Gu P, Xiao R, Pan G, et al. STCA: Spatio-Temporal Credit Assignment with Delayed Feedback in Deep Spiking Neural Networks[C]//IJCAI. 2019: 1366-1372. - .. [6] Roy D, Chakraborty I, Roy K. Scaling deep spiking neural networks with binary stochastic activations[C]//2019 IEEE International Conference on Cognitive Computing (ICCC). IEEE, 2019: 50-58. - .. [7] Cheng X, Hao Y, Xu J, et al. LISNN: Improving Spiking Neural Networks with Lateral Interactions for Robust Object Recognition[C]//IJCAI. 1519-1525. - .. [8] Kaiser J, Mostafa H, Neftci E. Synaptic plasticity dynamics for deep continuous local learning (DECOLLE)[J]. Frontiers in Neuroscience, 2020, 14: 424. - - """ - return PiecewiseLeakyRelu(c=c, w=w)(x) - - -class SquarewaveFourierSeries(_OneInpSurrogate): - """Judge spiking state with a squarewave fourier series. - - See Also:: - - squarewave_fourier_series - """ - - def __init__(self, n=2, t_period=8., forward_use_surrogate=False): - super().__init__(forward_use_surrogate=forward_use_surrogate) - self.n = n - self.t_period = t_period - - def surrogate_grad(self, dz, x): - x = as_jax(x) - w = jnp.pi * 2. / self.t_period - dx = jnp.cos(w * x) - for i in range(2, self.n): - dx += jnp.cos((2 * i - 1.) * w * x) - dx *= 4. / self.t_period - return dx * as_jax(dz) - - def surrogate_fun(self, x): - x = as_jax(x) - w = jnp.pi * 2. / self.t_period - ret = jnp.sin(w * x) - for i in range(2, self.n): - c = (2 * i - 1.) - ret += jnp.sin(c * w * x) / c - z = 0.5 + 2. / jnp.pi * ret - return z - - def __repr__(self): - return f'{self.__class__.__name__}(n={self.n}, t_period={self.t_period})' - - -def squarewave_fourier_series( - x: Union[jax.Array, Array], - n: int = 2, - t_period: float = 8., - origin: bool = False -): - r"""Judge spiking state with a squarewave fourier series. - - If `origin=False`, computes the forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - If `origin=True`, computes the original function: - - .. math:: - - g(x) = 0.5 + \frac{1}{\pi}*\sum_{i=1}^n {\sin\left({(2i-1)*2\pi}*x/T\right) \over 2i-1 } - - Backward function: - - .. math:: - - g'(x) = \sum_{i=1}^n\frac{4\cos\left((2 * i - 1.) * 2\pi * x / T\right)}{T} - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> bp.visualize.get_figure(1, 1, 4, 6) - >>> xs = bm.linspace(-3, 3, 1000) - >>> for n in [2, 4, 8]: - >>> f = bm.surrogate.SquarewaveFourierSeries(n=n) - >>> grads1 = bm.vector_grad(f)(xs) - >>> plt.plot(bm.as_numpy(xs), bm.as_numpy(grads1), label=f'n={n}') - >>> plt.legend() - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - n: int - t_period: float - origin: bool - Whether to compute the original function as the feedfoward output. - - Returns:: - - out: jax.Array - The spiking state. - - """ - - return SquarewaveFourierSeries(n=n, t_period=t_period, forward_use_surrogate=origin)(x) - - -class S2NN(_OneInpSurrogate): - """Judge spiking state with the S2NN surrogate spiking function. - - See Also:: - - s2nn - """ - - def __init__(self, alpha=4., beta=1., epsilon=1e-8, forward_use_surrogate=False): - super().__init__(forward_use_surrogate=forward_use_surrogate) - self.alpha = alpha - self.beta = beta - self.epsilon = epsilon - - def surrogate_fun(self, x): - x = as_jax(x) - z = jnp.where(x < 0., - sci.special.expit(x * self.alpha), - self.beta * jnp.log(jnp.abs((x + 1.)) + self.epsilon) + 0.5) - return z - - def surrogate_grad(self, dz, x): - x = as_jax(x) - sg = sci.special.expit(self.alpha * x) - dx = jnp.where(x < 0., self.alpha * sg * (1. - sg), self.beta / (x + 1.)) - return dx * as_jax(dz) - - def __repr__(self): - return f'{self.__class__.__name__}(alpha={self.alpha}, beta={self.beta}, epsilon={self.epsilon})' - - -def s2nn( - x: Union[jax.Array, Array], - alpha: float = 4., - beta: float = 1., - epsilon: float = 1e-8, - origin: bool = False -): - r"""Judge spiking state with the S2NN surrogate spiking function [1]_. - - If `origin=False`, computes the forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - If `origin=True`, computes the original function: - - .. math:: - - \begin{split}g(x) = \begin{cases} - \mathrm{sigmoid} (\alpha x), x < 0 \\ - \beta \ln(|x + 1|) + 0.5, x \ge 0 - \end{cases}\end{split} - - Backward function: - - .. math:: - - \begin{split}g'(x) = \begin{cases} - \alpha * (1 - \mathrm{sigmoid} (\alpha x)) \mathrm{sigmoid} (\alpha x), x < 0 \\ - \frac{\beta}{(x + 1)}, x \ge 0 - \end{cases}\end{split} - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> bp.visualize.get_figure(1, 1, 4, 6) - >>> xs = bm.linspace(-3, 3, 1000) - >>> grads = bm.vector_grad(bm.surrogate.s2nn)(xs, 4., 1.) - >>> plt.plot(bm.as_numpy(xs), bm.as_numpy(grads), label=r'$\alpha=4, \beta=1$') - >>> grads = bm.vector_grad(bm.surrogate.s2nn)(xs, 8., 2.) - >>> plt.plot(bm.as_numpy(xs), bm.as_numpy(grads), label=r'$\alpha=8, \beta=2$') - >>> plt.legend() - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - alpha: float - The param that controls the gradient when ``x < 0``. - beta: float - The param that controls the gradient when ``x >= 0`` - epsilon: float - Avoid nan - origin: bool - Whether to compute the original function as the feedfoward output. - - Returns:: - - out: jax.Array - The spiking state. - - References:: - - .. [1] Suetake, Kazuma et al. “S2NN: Time Step Reduction of Spiking Surrogate Gradients for Training Energy Efficient Single-Step Neural Networks.” ArXiv abs/2201.10879 (2022): n. pag. - - """ - return S2NN(alpha=alpha, beta=beta, epsilon=epsilon, forward_use_surrogate=origin)(x) - - -class QPseudoSpike(_OneInpSurrogate): - """Judge spiking state with the q-PseudoSpike surrogate function. - - See Also:: - - q_pseudo_spike - """ - - def __init__(self, alpha=2., forward_use_surrogate=False): - super().__init__(forward_use_surrogate=forward_use_surrogate) - self.alpha = alpha - - def surrogate_grad(self, dz, x): - x = as_jax(x) - dx = jnp.power(1 + 2 / (self.alpha - 1) * jnp.abs(x), -self.alpha) - return dx * as_jax(dz) - - def surrogate_fun(self, x): - x = as_jax(x) - z = jnp.where(x < 0., - 0.5 * jnp.power(1 - 2 / (self.alpha - 1) * x, 1 - self.alpha), - 1. - 0.5 * jnp.power(1 + 2 / (self.alpha - 1) * jnp.abs(x), 1 - self.alpha)) - return z - - def __repr__(self): - return f'{self.__class__.__name__}(alpha={self.alpha})' - - -def q_pseudo_spike( - x: Union[jax.Array, Array], - alpha: float = 2., - origin: bool = False -): - r"""Judge spiking state with the q-PseudoSpike surrogate function [1]_. - - If `origin=False`, computes the forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - If `origin=True`, computes the original function: - - .. math:: - - \begin{split}g(x) = - \begin{cases} - \frac{1}{2}(1-\frac{2x}{\alpha-1})^{1-\alpha}, & x < 0 \\ - 1 - \frac{1}{2}(1+\frac{2x}{\alpha-1})^{1-\alpha}, & x \geq 0. - \end{cases}\end{split} - - Backward function: - - .. math:: - - g'(x) = (1+\frac{2|x|}{\alpha-1})^{-\alpha} - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> xs = bm.linspace(-3, 3, 1000) - >>> bp.visualize.get_figure(1, 1, 4, 6) - >>> for alpha in [0.5, 1., 2., 4.]: - >>> grads = bm.vector_grad(bm.surrogate.q_pseudo_spike)(xs, alpha) - >>> plt.plot(bm.as_numpy(xs), bm.as_numpy(grads), label=r'$\alpha=$' + str(alpha)) - >>> plt.legend() - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - alpha: float - The parameter to control tail fatness of gradient. - origin: bool - Whether to compute the original function as the feedfoward output. - - Returns:: - - out: jax.Array - The spiking state. - - References:: - - .. [1] Herranz-Celotti, Luca and Jean Rouat. “Surrogate Gradients Design.” ArXiv abs/2202.00282 (2022): n. pag. - """ - return QPseudoSpike(alpha=alpha, forward_use_surrogate=origin)(x) - - -class LeakyRelu(_OneInpSurrogate): - """Judge spiking state with the Leaky ReLU function. - - See Also:: - - leaky_relu - """ - - def __init__(self, alpha=0.1, beta=1., forward_use_surrogate=False): - super().__init__(forward_use_surrogate=forward_use_surrogate) - self.alpha = alpha - self.beta = beta - - def surrogate_fun(self, x): - x = as_jax(x) - return jnp.where(x < 0., self.alpha * x, self.beta * x) - - def surrogate_grad(self, dz, x): - x = as_jax(x) - dx = jnp.where(x < 0., self.alpha, self.beta) - return dx * as_jax(dz) - - def __repr__(self): - return f'{self.__class__.__name__}(alpha={self.alpha}, beta={self.beta})' - - -def leaky_relu( - x: Union[jax.Array, Array], - alpha: float = 0.1, - beta: float = 1., - origin: bool = False -): - r"""Judge spiking state with the Leaky ReLU function. - - If `origin=False`, computes the forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - If `origin=True`, computes the original function: - - .. math:: - - \begin{split}g(x) = - \begin{cases} - \beta \cdot x, & x \geq 0 \\ - \alpha \cdot x, & x < 0 \\ - \end{cases}\end{split} - - Backward function: - - .. math:: - - \begin{split}g'(x) = - \begin{cases} - \beta, & x \geq 0 \\ - \alpha, & x < 0 \\ - \end{cases}\end{split} - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> xs = bm.linspace(-3, 3, 1000) - >>> bp.visualize.get_figure(1, 1, 4, 6) - >>> grads = bm.vector_grad(bm.surrogate.leaky_relu)(xs, 0., 1.) - >>> plt.plot(bm.as_numpy(xs), bm.as_numpy(grads), label=r'$\alpha=0., \beta=1.$') - >>> plt.legend() - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - alpha: float - The parameter to control the gradient when :math:`x < 0`. - beta: float - The parameter to control the gradient when :math:`x >= 0`. - origin: bool - Whether to compute the original function as the feedfoward output. - - Returns:: - - out: jax.Array - The spiking state. - """ - return LeakyRelu(alpha=alpha, beta=beta, forward_use_surrogate=origin)(x) - - -class LogTailedRelu(_OneInpSurrogate): - """Judge spiking state with the Log-tailed ReLU function. - - See Also:: - - log_tailed_relu - """ - - def __init__(self, alpha=0., forward_use_surrogate=False): - super().__init__(forward_use_surrogate=forward_use_surrogate) - self.alpha = alpha - - def surrogate_fun(self, x): - x = as_jax(x) - z = jnp.where(x > 1, - jnp.log(x), - jnp.where(x > 0, - x, - self.alpha * x)) - return z - - def surrogate_grad(self, dz, x): - x = as_jax(x) - dx = jnp.where(x > 1, - 1 / x, - jnp.where(x > 0, - 1., - self.alpha)) - return dx * as_jax(dz) - - def __repr__(self): - return f'{self.__class__.__name__}(alpha={self.alpha})' - - -def log_tailed_relu( - x: Union[jax.Array, Array], - alpha: float = 0., - origin: bool = False -): - r"""Judge spiking state with the Log-tailed ReLU function [1]_. - - If `origin=False`, computes the forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - If `origin=True`, computes the original function: - - .. math:: - - \begin{split}g(x) = - \begin{cases} - \alpha x, & x \leq 0 \\ - x, & 0 < x \leq 0 \\ - log(x), x > 1 \\ - \end{cases}\end{split} - - Backward function: - - .. math:: - - \begin{split}g'(x) = - \begin{cases} - \alpha, & x \leq 0 \\ - 1, & 0 < x \leq 0 \\ - \frac{1}{x}, x > 1 \\ - \end{cases}\end{split} - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> xs = bm.linspace(-3, 3, 1000) - >>> bp.visualize.get_figure(1, 1, 4, 6) - >>> grads = bm.vector_grad(bm.surrogate.leaky_relu)(xs, 0., 1.) - >>> plt.plot(bm.as_numpy(xs), bm.as_numpy(grads), label=r'$\alpha=0., \beta=1.$') - >>> plt.legend() - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - alpha: float - The parameter to control the gradient. - origin: bool - Whether to compute the original function as the feedfoward output. - - Returns:: - - out: jax.Array - The spiking state. - - References:: - - .. [1] Cai, Zhaowei et al. “Deep Learning with Low Precision by Half-Wave Gaussian Quantization.” 2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR) (2017): 5406-5414. - """ - return LogTailedRelu(alpha=alpha, forward_use_surrogate=origin)(x) - - -class ReluGrad(_OneInpSurrogate): - """Judge spiking state with the ReLU gradient function. - - See Also:: - - relu_grad - """ - - def __init__(self, alpha=0.3, width=1.): - super().__init__() - self.alpha = alpha - self.width = width - - def surrogate_grad(self, dz, x): - x = as_jax(x) - dx = jnp.maximum(self.alpha * self.width - jnp.abs(x) * self.alpha, 0) - return dx * as_jax(dz) - - def __repr__(self): - return f'{self.__class__.__name__}(alpha={self.alpha}, width={self.width})' - - -def relu_grad( - x: Union[jax.Array, Array], - alpha: float = 0.3, - width: float = 1., -): - r"""Spike function with the ReLU gradient function [1]_. - - The forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - Backward function: - - .. math:: - - g'(x) = \text{ReLU}(\alpha * (\mathrm{width}-|x|)) - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> xs = bm.linspace(-3, 3, 1000) - >>> bp.visualize.get_figure(1, 1, 4, 6) - >>> for s in [0.5, 1.]: - >>> for w in [1, 2.]: - >>> grads = bm.vector_grad(bm.surrogate.relu_grad)(xs, s, w) - >>> plt.plot(bm.as_numpy(xs), bm.as_numpy(grads), label=r'$\alpha=$' + f'{s}, width={w}') - >>> plt.legend() - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - alpha: float - The parameter to control the gradient. - width: float - The parameter to control the width of the gradient. - - Returns:: - - out: jax.Array - The spiking state. - - References:: - - .. [1] Neftci, E. O., Mostafa, H. & Zenke, F. Surrogate gradient learning in spiking neural networks. IEEE Signal Process. Mag. 36, 61–63 (2019). - """ - return ReluGrad(alpha=alpha, width=width)(x) - - -class GaussianGrad(_OneInpSurrogate): - """Judge spiking state with the Gaussian gradient function. - - See Also:: - - gaussian_grad - """ - - def __init__(self, sigma=0.5, alpha=0.5): - super().__init__() - self.sigma = sigma - self.alpha = alpha - - def surrogate_grad(self, dz, x): - x = as_jax(x) - dx = jnp.exp(-(x ** 2) / (2 * jnp.power(self.sigma, 2))) / (jnp.sqrt(2 * jnp.pi) * self.sigma) - return self.alpha * dx * as_jax(dz) - - def __repr__(self): - return f'{self.__class__.__name__}(alpha={self.alpha}, sigma={self.sigma})' - - -def gaussian_grad( - x: Union[jax.Array, Array], - sigma: float = 0.5, - alpha: float = 0.5, -): - r"""Spike function with the Gaussian gradient function [1]_. - - The forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - Backward function: - - .. math:: - - g'(x) = \alpha * \text{gaussian}(x, 0., \sigma) - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> xs = bm.linspace(-3, 3, 1000) - >>> bp.visualize.get_figure(1, 1, 4, 6) - >>> for s in [0.5, 1., 2.]: - >>> grads = bm.vector_grad(bm.surrogate.gaussian_grad)(xs, s, 0.5) - >>> plt.plot(bm.as_numpy(xs), bm.as_numpy(grads), label=r'$\alpha=0.5, \sigma=$' + str(s)) - >>> plt.legend() - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - sigma: float - The parameter to control the variance of gaussian distribution. - alpha: float - The parameter to control the scale of the gradient. - - Returns:: - - out: jax.Array - The spiking state. - - References:: - - .. [1] Yin, B., Corradi, F. & Bohté, S.M. Accurate and efficient time-domain classification with adaptive spiking recurrent neural networks. Nat Mach Intell 3, 905–913 (2021). - """ - return GaussianGrad(sigma=sigma, alpha=alpha)(x) - - -class MultiGaussianGrad(_OneInpSurrogate): - """Judge spiking state with the multi-Gaussian gradient function. - - See Also:: - - multi_gaussian_grad - """ - - def __init__(self, h=0.15, s=6.0, sigma=0.5, scale=0.5): - super().__init__() - self.h = h - self.s = s - self.sigma = sigma - self.scale = scale - - def surrogate_grad(self, dz, x): - x = as_jax(x) - g1 = jnp.exp(-x ** 2 / (2 * jnp.power(self.sigma, 2))) / (jnp.sqrt(2 * jnp.pi) * self.sigma) - g2 = jnp.exp(-(x - self.sigma) ** 2 / (2 * jnp.power(self.s * self.sigma, 2)) - ) / (jnp.sqrt(2 * jnp.pi) * self.s * self.sigma) - g3 = jnp.exp(-(x + self.sigma) ** 2 / (2 * jnp.power(self.s * self.sigma, 2)) - ) / (jnp.sqrt(2 * jnp.pi) * self.s * self.sigma) - dx = g1 * (1. + self.h) - g2 * self.h - g3 * self.h - return self.scale * dx * as_jax(dz) - - def __repr__(self): - return f'{self.__class__.__name__}(h={self.h}, s={self.s}, sigma={self.sigma}, scale={self.scale})' - - -def multi_gaussian_grad( - x: Union[jax.Array, Array], - h: float = 0.15, - s: float = 6.0, - sigma: float = 0.5, - scale: float = 0.5, -): - r"""Spike function with the multi-Gaussian gradient function [1]_. - - The forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - Backward function: - - .. math:: - - \begin{array}{l} - g'(x)=(1+h){{{\mathcal{N}}}}(x, 0, {\sigma }^{2}) - -h{{{\mathcal{N}}}}(x, \sigma,{(s\sigma )}^{2})- - h{{{\mathcal{N}}}}(x, -\sigma ,{(s\sigma )}^{2}) - \end{array} - - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> xs = bm.linspace(-3, 3, 1000) - >>> bp.visualize.get_figure(1, 1, 4, 6) - >>> grads = bm.vector_grad(bm.surrogate.multi_gaussian_grad)(xs) - >>> plt.plot(bm.as_numpy(xs), bm.as_numpy(grads)) - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - h: float - The hyper-parameters of approximate function - s: float - The hyper-parameters of approximate function - sigma: float - The gaussian sigma. - scale: float - The gradient scale. - - Returns:: - - out: jax.Array - The spiking state. - - References:: - - .. [1] Yin, B., Corradi, F. & Bohté, S.M. Accurate and efficient time-domain classification with adaptive spiking recurrent neural networks. Nat Mach Intell 3, 905–913 (2021). - """ - return MultiGaussianGrad(h=h, s=s, sigma=sigma, scale=scale)(x) - - -class InvSquareGrad(_OneInpSurrogate): - """Judge spiking state with the inverse-square surrogate gradient function. - - See Also:: - - inv_square_grad - """ - - def __init__(self, alpha=100.): - super().__init__() - self.alpha = alpha - - def surrogate_grad(self, dz, x): - x = as_jax(x) - dx = as_jax(dz) / (self.alpha * jnp.abs(x) + 1.0) ** 2 - return dx - - def __repr__(self): - return f'{self.__class__.__name__}(alpha={self.alpha})' - - -def inv_square_grad( - x: Union[jax.Array, Array], - alpha: float = 100. -): - r"""Spike function with the inverse-square surrogate gradient. - - Forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - Backward function: - - .. math:: - - g'(x) = \frac{1}{(\alpha * |x| + 1.) ^ 2} - - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> xs = bm.linspace(-1, 1, 1000) - >>> for alpha in [1., 10., 100.]: - >>> grads = bm.vector_grad(bm.surrogate.inv_square_grad)(xs, alpha) - >>> plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha)) - >>> plt.legend() - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - alpha: float - Parameter to control smoothness of gradient - - Returns:: - - out: jax.Array - The spiking state. - """ - return InvSquareGrad(alpha=alpha)(x) - - -class SlayerGrad(_OneInpSurrogate): - """Judge spiking state with the slayer surrogate gradient function. - - See Also:: - - slayer_grad - """ - - def __init__(self, alpha=1.): - super().__init__() - self.alpha = alpha - - def surrogate_grad(self, dz, x): - x = as_jax(x) - dx = as_jax(dz) * jnp.exp(-self.alpha * jnp.abs(x)) - return dx - - def __repr__(self): - return f'{self.__class__.__name__}(alpha={self.alpha})' - - -def slayer_grad( - x: Union[jax.Array, Array], - alpha: float = 1. -): - r"""Spike function with the slayer surrogate gradient function. - - Forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - Backward function: - - .. math:: - - g'(x) = \exp(-\alpha |x|) - - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> bp.visualize.get_figure(1, 1, 4, 6) - >>> xs = bm.linspace(-3, 3, 1000) - >>> for alpha in [0.5, 1., 2., 4.]: - >>> grads = bm.vector_grad(bm.surrogate.slayer_grad)(xs, alpha) - >>> plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha)) - >>> plt.legend() - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - alpha: float - Parameter to control smoothness of gradient - - Returns:: - - out: jax.Array - The spiking state. - - References:: - - .. [1] Shrestha, S. B. & Orchard, G. Slayer: spike layer error reassignment in time. In Advances in Neural Information Processing Systems Vol. 31, 1412–1421 (NeurIPS, 2018). - """ - return SlayerGrad(alpha=alpha)(x) diff --git a/brainpy/math/surrogate/_one_input_new.py b/brainpy/math/surrogate/_one_input_new.py deleted file mode 100644 index 908ab5eb8..000000000 --- a/brainpy/math/surrogate/_one_input_new.py +++ /dev/null @@ -1,1684 +0,0 @@ -# -*- coding: utf-8 -*- -# Copyright 2025 BrainX 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 Union - -import jax -import jax.numpy as jnp -import jax.scipy as sci -from brainstate._compatible_import import Primitive -from jax.interpreters import batching, ad, mlir - -from brainpy.math.interoperability import as_jax -from brainpy.math.ndarray import Array as Array - -__all__ = [ - 'Surrogate', - 'Sigmoid', - 'sigmoid', - 'PiecewiseQuadratic', - 'piecewise_quadratic', - 'PiecewiseExp', - 'piecewise_exp', - 'SoftSign', - 'soft_sign', - 'Arctan', - 'arctan', - 'NonzeroSignLog', - 'nonzero_sign_log', - 'ERF', - 'erf', - 'PiecewiseLeakyRelu', - 'piecewise_leaky_relu', - 'SquarewaveFourierSeries', - 'squarewave_fourier_series', - 'S2NN', - 's2nn', - 'QPseudoSpike', - 'q_pseudo_spike', - 'LeakyRelu', - 'leaky_relu', - 'LogTailedRelu', - 'log_tailed_relu', - 'ReluGrad', - 'relu_grad', - 'GaussianGrad', - 'gaussian_grad', - 'InvSquareGrad', - 'inv_square_grad', - 'MultiGaussianGrad', - 'multi_gaussian_grad', - 'SlayerGrad', - 'slayer_grad', -] - - -def _heaviside_abstract(x, dx): - return [x] - - -def _heaviside_imp(x, dx): - z = jnp.asarray(x >= 0, dtype=x.dtype) - return [z] - - -def _heaviside_batching(args, axes): - return heaviside_p.bind(*args), [axes[0]] - - -def _heaviside_jvp(primals, tangents): - x, dx = primals - tx, tdx = tangents - primal_outs = heaviside_p.bind(x, dx) - tangent_outs = [dx * tx, ] - return primal_outs, tangent_outs - - -heaviside_p = Primitive('heaviside_p') -heaviside_p.multiple_results = True -heaviside_p.def_abstract_eval(_heaviside_abstract) -heaviside_p.def_impl(_heaviside_imp) -batching.primitive_batchers[heaviside_p] = _heaviside_batching -ad.primitive_jvps[heaviside_p] = _heaviside_jvp -mlir.register_lowering(heaviside_p, mlir.lower_fun(_heaviside_imp, multiple_results=True)) - - -def _is_bp_array(x): - return isinstance(x, Array) - - -def _as_jax(x): - return x.value if _is_bp_array(x) else x - - -class Surrogate(object): - """The base surrograte gradient function. - - To customize a surrogate gradient function, you can inherit this class and - implement the `surrogate_fun` and `surrogate_grad` methods. - - Examples:: - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import jax.numpy as jnp - - >>> class MySurrogate(bm.Surrogate): - ... def __init__(self, alpha=1.): - ... super().__init__() - ... self.alpha = alpha - ... - ... def surrogate_fun(self, x): - ... return jnp.sin(x) * self.alpha - ... - ... def surrogate_grad(self, x): - ... return jnp.cos(x) * self.alpha - - """ - - def __call__(self, x): - x = _as_jax(x) - dx = self.surrogate_grad(x) - return heaviside_p.bind(x, dx)[0] - - def __repr__(self): - return f'{self.__class__.__name__}()' - - def surrogate_fun(self, x) -> jax.Array: - """The surrogate function.""" - raise NotImplementedError - - def surrogate_grad(self, x) -> jax.Array: - """The gradient function of the surrogate function.""" - raise NotImplementedError - - -class Sigmoid(Surrogate): - """Spike function with the sigmoid-shaped surrogate gradient. - - See Also:: - - sigmoid - - """ - - def __init__(self, alpha: float = 4.): - super().__init__() - self.alpha = alpha - - def surrogate_fun(self, x): - return sci.special.expit(self.alpha * x) - - def surrogate_grad(self, x): - sgax = sci.special.expit(x * self.alpha) - dx = (1. - sgax) * sgax * self.alpha - return dx - - def __repr__(self): - return f'{self.__class__.__name__}(alpha={self.alpha})' - - -def sigmoid( - x: Union[jax.Array, Array], - alpha: float = 4., -): - r"""Spike function with the sigmoid-shaped surrogate gradient. - - The forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - - Backward function: - - .. math:: - - g'(x) = \alpha * (1 - \mathrm{sigmoid} (\alpha x)) \mathrm{sigmoid} (\alpha x) - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> bp.visualize.get_figure(1, 1, 4, 6) - >>> xs = bm.linspace(-2, 2, 1000) - >>> for alpha in [1., 2., 4.]: - >>> grads = bm.vector_grad(bm.surrogate.sigmoid)(xs, alpha) - >>> plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha)) - >>> plt.legend() - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - alpha: float - Parameter to control smoothness of gradient - - - Returns:: - - out: jax.Array - The spiking state. - """ - return Sigmoid(alpha=alpha)(x) - - -class PiecewiseQuadratic(Surrogate): - """Judge spiking state with a piecewise quadratic function. - - See Also:: - - piecewise_quadratic - - """ - - def __init__(self, alpha: float = 1.): - super().__init__() - self.alpha = alpha - - def surrogate_fun(self, x): - x = as_jax(x) - z = jnp.where(x < -1 / self.alpha, - 0., - jnp.where(x > 1 / self.alpha, - 1., - (-self.alpha * jnp.abs(x) / 2 + 1) * self.alpha * x + 0.5)) - return z - - def surrogate_grad(self, x): - x = as_jax(x) - dx = jnp.where(jnp.abs(x) > 1 / self.alpha, 0., (-self.alpha ** 2 * jnp.abs(x) + self.alpha)) - return dx - - def __repr__(self): - return f'{self.__class__.__name__}(alpha={self.alpha})' - - -def piecewise_quadratic( - x: Union[jax.Array, Array], - alpha: float = 1., -): - r"""Judge spiking state with a piecewise quadratic function [1]_ [2]_ [3]_ [4]_ [5]_. - - The forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - - Backward function: - - .. math:: - - g'(x) = - \begin{cases} - 0, & |x| > \frac{1}{\alpha} \\ - -\alpha^2|x|+\alpha, & |x| \leq \frac{1}{\alpha} - \end{cases} - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> bp.visualize.get_figure(1, 1, 4, 6) - >>> xs = bm.linspace(-3, 3, 1000) - >>> for alpha in [0.5, 1., 2., 4.]: - >>> grads = bm.vector_grad(bm.surrogate.piecewise_quadratic)(xs, alpha) - >>> plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha)) - >>> plt.legend() - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - alpha: float - Parameter to control smoothness of gradient - - - Returns:: - - out: jax.Array - The spiking state. - - References:: - - .. [1] Esser S K, Merolla P A, Arthur J V, et al. Convolutional networks for fast, energy-efficient neuromorphic computing[J]. Proceedings of the national academy of sciences, 2016, 113(41): 11441-11446. - .. [2] Wu Y, Deng L, Li G, et al. Spatio-temporal backpropagation for training high-performance spiking neural networks[J]. Frontiers in neuroscience, 2018, 12: 331. - .. [3] Bellec G, Salaj D, Subramoney A, et al. Long short-term memory and learning-to-learn in networks of spiking neurons[C]//Proceedings of the 32nd International Conference on Neural Information Processing Systems. 2018: 795-805. - .. [4] Neftci E O, Mostafa H, Zenke F. Surrogate gradient learning in spiking neural networks: Bringing the power of gradient-based optimization to spiking neural networks[J]. IEEE Signal Processing Magazine, 2019, 36(6): 51-63. - .. [5] Panda P, Aketi S A, Roy K. Toward scalable, efficient, and accurate deep spiking neural networks with backward residual connections, stochastic softmax, and hybridization[J]. Frontiers in Neuroscience, 2020, 14. - """ - return PiecewiseQuadratic(alpha=alpha)(x) - - -class PiecewiseExp(Surrogate): - """Judge spiking state with a piecewise exponential function. - - See Also:: - - piecewise_exp - """ - - def __init__(self, alpha: float = 1.): - super().__init__() - self.alpha = alpha - - def surrogate_grad(self, x): - x = as_jax(x) - dx = (self.alpha / 2) * jnp.exp(-self.alpha * jnp.abs(x)) - return dx - - def surrogate_fun(self, x): - x = as_jax(x) - return jnp.where(x < 0, jnp.exp(self.alpha * x) / 2, 1 - jnp.exp(-self.alpha * x) / 2) - - def __repr__(self): - return f'{self.__class__.__name__}(alpha={self.alpha})' - - -def piecewise_exp( - x: Union[jax.Array, Array], - alpha: float = 1., - -): - r"""Judge spiking state with a piecewise exponential function [1]_. - - The forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - - Backward function: - - .. math:: - - g'(x) = \frac{\alpha}{2}e^{-\alpha |x|} - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> bp.visualize.get_figure(1, 1, 4, 6) - >>> xs = bm.linspace(-3, 3, 1000) - >>> for alpha in [0.5, 1., 2., 4.]: - >>> grads = bm.vector_grad(bm.surrogate.piecewise_exp)(xs, alpha) - >>> plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha)) - >>> plt.legend() - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - alpha: float - Parameter to control smoothness of gradient - - - Returns:: - - out: jax.Array - The spiking state. - - References:: - - .. [1] Neftci E O, Mostafa H, Zenke F. Surrogate gradient learning in spiking neural networks: Bringing the power of gradient-based optimization to spiking neural networks[J]. IEEE Signal Processing Magazine, 2019, 36(6): 51-63. - """ - return PiecewiseExp(alpha=alpha)(x) - - -class SoftSign(Surrogate): - """Judge spiking state with a soft sign function. - - See Also:: - - soft_sign - """ - - def __init__(self, alpha=1.): - super().__init__() - self.alpha = alpha - - def surrogate_grad(self, x): - x = as_jax(x) - dx = self.alpha * 0.5 / (1 + jnp.abs(self.alpha * x)) ** 2 - return dx - - def surrogate_fun(self, x): - x = as_jax(x) - return x / (2 / self.alpha + 2 * jnp.abs(x)) + 0.5 - - def __repr__(self): - return f'{self.__class__.__name__}(alpha={self.alpha})' - - -def soft_sign( - x: Union[jax.Array, Array], - alpha: float = 1., - -): - r"""Judge spiking state with a soft sign function. - - The forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - - Backward function: - - .. math:: - - g'(x) = \frac{\alpha}{2(1 + |\alpha x|)^{2}} = \frac{1}{2\alpha(\frac{1}{\alpha} + |x|)^{2}} - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> bp.visualize.get_figure(1, 1, 4, 6) - >>> xs = bm.linspace(-3, 3, 1000) - >>> for alpha in [0.5, 1., 2., 4.]: - >>> grads = bm.vector_grad(bm.surrogate.soft_sign)(xs, alpha) - >>> plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha)) - >>> plt.legend() - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - alpha: float - Parameter to control smoothness of gradient - - - Returns:: - - out: jax.Array - The spiking state. - - """ - return SoftSign(alpha=alpha)(x) - - -class Arctan(Surrogate): - """Judge spiking state with an arctan function. - - See Also:: - - arctan - """ - - def __init__(self, alpha=1.): - super().__init__() - self.alpha = alpha - - def surrogate_grad(self, x): - x = as_jax(x) - dx = self.alpha * 0.5 / (1 + (jnp.pi / 2 * self.alpha * x) ** 2) - return dx - - def surrogate_fun(self, x): - x = as_jax(x) - return jnp.arctan(jnp.pi / 2 * self.alpha * x) / jnp.pi + 0.5 - - def __repr__(self): - return f'{self.__class__.__name__}(alpha={self.alpha})' - - -def arctan( - x: Union[jax.Array, Array], - alpha: float = 1., - -): - r"""Judge spiking state with an arctan function. - - The forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - - Backward function: - - .. math:: - - g'(x) = \frac{\alpha}{2(1 + (\frac{\pi}{2}\alpha x)^2)} - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> bp.visualize.get_figure(1, 1, 4, 6) - >>> xs = bm.linspace(-3, 3, 1000) - >>> for alpha in [0.5, 1., 2., 4.]: - >>> grads = bm.vector_grad(bm.surrogate.arctan)(xs, alpha) - >>> plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha)) - >>> plt.legend() - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - alpha: float - Parameter to control smoothness of gradient - - - Returns:: - - out: jax.Array - The spiking state. - - """ - return Arctan(alpha=alpha)(x) - - -class NonzeroSignLog(Surrogate): - """Judge spiking state with a nonzero sign log function. - - See Also:: - - nonzero_sign_log - """ - - def __init__(self, alpha=1.): - super().__init__() - self.alpha = alpha - - def surrogate_grad(self, x): - x = as_jax(x) - dx = 1. / (1 / self.alpha + jnp.abs(x)) - return dx - - def surrogate_fun(self, x): - x = as_jax(x) - return jnp.where(x < 0, -1., 1.) * jnp.log(jnp.abs(self.alpha * x) + 1) - - def __repr__(self): - return f'{self.__class__.__name__}(alpha={self.alpha})' - - -def nonzero_sign_log( - x: Union[jax.Array, Array], - alpha: float = 1., - -): - r"""Judge spiking state with a nonzero sign log function. - - The forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - - Backward function: - - .. math:: - - g'(x) = \frac{\alpha}{1 + |\alpha x|} = \frac{1}{\frac{1}{\alpha} + |x|} - - This surrogate function has the advantage of low computation cost during the backward. - - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> bp.visualize.get_figure(1, 1, 4, 6) - >>> xs = bm.linspace(-3, 3, 1000) - >>> for alpha in [0.5, 1., 2., 4.]: - >>> grads = bm.vector_grad(bm.surrogate.nonzero_sign_log)(xs, alpha) - >>> plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha)) - >>> plt.legend() - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - alpha: float - Parameter to control smoothness of gradient - - - Returns:: - - out: jax.Array - The spiking state. - - """ - return NonzeroSignLog(alpha=alpha)(x) - - -class ERF(Surrogate): - """Judge spiking state with an erf function. - - See Also:: - - erf - """ - - def __init__(self, alpha=1.): - super().__init__() - self.alpha = alpha - - def surrogate_grad(self, x): - x = as_jax(x) - dx = (self.alpha / jnp.sqrt(jnp.pi)) * jnp.exp(-jnp.power(self.alpha, 2) * x * x) - return dx - - def surrogate_fun(self, x): - x = as_jax(x) - return 0.5 * (1. - sci.special.erf(-self.alpha * x)) - - def __repr__(self): - return f'{self.__class__.__name__}(alpha={self.alpha})' - - -def erf( - x: Union[jax.Array, Array], - alpha: float = 1., - -): - r"""Judge spiking state with an erf function [1]_ [2]_ [3]_. - - The forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - - Backward function: - - .. math:: - - g'(x) = \frac{\alpha}{\sqrt{\pi}}e^{-\alpha^2x^2} - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> bp.visualize.get_figure(1, 1, 4, 6) - >>> xs = bm.linspace(-3, 3, 1000) - >>> for alpha in [0.5, 1., 2., 4.]: - >>> grads = bm.vector_grad(bm.surrogate.nonzero_sign_log)(xs, alpha) - >>> plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha)) - >>> plt.legend() - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - alpha: float - Parameter to control smoothness of gradient - - - Returns:: - - out: jax.Array - The spiking state. - - References:: - - .. [1] Esser S K, Appuswamy R, Merolla P, et al. Backpropagation for energy-efficient neuromorphic computing[J]. Advances in neural information processing systems, 2015, 28: 1117-1125. - .. [2] Wu Y, Deng L, Li G, et al. Spatio-temporal backpropagation for training high-performance spiking neural networks[J]. Frontiers in neuroscience, 2018, 12: 331. - .. [3] Yin B, Corradi F, Bohté S M. Effective and efficient computation with multiple-timescale spiking recurrent neural networks[C]//International Conference on Neuromorphic Systems 2020. 2020: 1-8. - - """ - return ERF(alpha=alpha)(x) - - -class PiecewiseLeakyRelu(Surrogate): - """Judge spiking state with a piecewise leaky relu function. - - See Also:: - - piecewise_leaky_relu - """ - - def __init__(self, c=0.01, w=1.): - super().__init__() - self.c = c - self.w = w - - def surrogate_fun(self, x): - x = as_jax(x) - z = jnp.where(x < -self.w, - self.c * x + self.c * self.w, - jnp.where(x > self.w, - self.c * x - self.c * self.w + 1, - 0.5 * x / self.w + 0.5)) - return z - - def surrogate_grad(self, x): - x = as_jax(x) - dx = jnp.where(jnp.abs(x) > self.w, self.c, 1 / self.w) - return dx - - def __repr__(self): - return f'{self.__class__.__name__}(c={self.c}, w={self.w})' - - -def piecewise_leaky_relu( - x: Union[jax.Array, Array], - c: float = 0.01, - w: float = 1., - -): - r"""Judge spiking state with a piecewise leaky relu function [1]_ [2]_ [3]_ [4]_ [5]_ [6]_ [7]_ [8]_. - - The forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - - Backward function: - - .. math:: - - \begin{split}g'(x) = - \begin{cases} - \frac{1}{w}, & |x| \leq w \\ - c, & |x| > w - \end{cases}\end{split} - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> bp.visualize.get_figure(1, 1, 4, 6) - >>> xs = bm.linspace(-3, 3, 1000) - >>> for c in [0.01, 0.05, 0.1]: - >>> for w in [1., 2.]: - >>> grads1 = bm.vector_grad(bm.surrogate.piecewise_leaky_relu)(xs, c=c, w=w) - >>> plt.plot(bm.as_numpy(xs), bm.as_numpy(grads1), label=f'x={c}, w={w}') - >>> plt.legend() - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - c: float - When :math:`|x| > w` the gradient is `c`. - w: float - When :math:`|x| <= w` the gradient is `1 / w`. - - - Returns:: - - out: jax.Array - The spiking state. - - References:: - - .. [1] Yin S, Venkataramanaiah S K, Chen G K, et al. Algorithm and hardware design of discrete-time spiking neural networks based on back propagation with binary activations[C]//2017 IEEE Biomedical Circuits and Systems Conference (BioCAS). IEEE, 2017: 1-5. - .. [2] Wu Y, Deng L, Li G, et al. Spatio-temporal backpropagation for training high-performance spiking neural networks[J]. Frontiers in neuroscience, 2018, 12: 331. - .. [3] Huh D, Sejnowski T J. Gradient descent for spiking neural networks[C]//Proceedings of the 32nd International Conference on Neural Information Processing Systems. 2018: 1440-1450. - .. [4] Wu Y, Deng L, Li G, et al. Direct training for spiking neural networks: Faster, larger, better[C]//Proceedings of the AAAI Conference on Artificial Intelligence. 2019, 33(01): 1311-1318. - .. [5] Gu P, Xiao R, Pan G, et al. STCA: Spatio-Temporal Credit Assignment with Delayed Feedback in Deep Spiking Neural Networks[C]//IJCAI. 2019: 1366-1372. - .. [6] Roy D, Chakraborty I, Roy K. Scaling deep spiking neural networks with binary stochastic activations[C]//2019 IEEE International Conference on Cognitive Computing (ICCC). IEEE, 2019: 50-58. - .. [7] Cheng X, Hao Y, Xu J, et al. LISNN: Improving Spiking Neural Networks with Lateral Interactions for Robust Object Recognition[C]//IJCAI. 1519-1525. - .. [8] Kaiser J, Mostafa H, Neftci E. Synaptic plasticity dynamics for deep continuous local learning (DECOLLE)[J]. Frontiers in Neuroscience, 2020, 14: 424. - - """ - return PiecewiseLeakyRelu(c=c, w=w)(x) - - -class SquarewaveFourierSeries(Surrogate): - """Judge spiking state with a squarewave fourier series. - - See Also:: - - squarewave_fourier_series - """ - - def __init__(self, n=2, t_period=8.): - super().__init__() - self.n = n - self.t_period = t_period - - def surrogate_grad(self, x): - x = as_jax(x) - w = jnp.pi * 2. / self.t_period - dx = jnp.cos(w * x) - for i in range(2, self.n): - dx += jnp.cos((2 * i - 1.) * w * x) - dx *= 4. / self.t_period - return dx - - def surrogate_fun(self, x): - x = as_jax(x) - w = jnp.pi * 2. / self.t_period - ret = jnp.sin(w * x) - for i in range(2, self.n): - c = (2 * i - 1.) - ret += jnp.sin(c * w * x) / c - z = 0.5 + 2. / jnp.pi * ret - return z - - def __repr__(self): - return f'{self.__class__.__name__}(n={self.n}, t_period={self.t_period})' - - -def squarewave_fourier_series( - x: Union[jax.Array, Array], - n: int = 2, - t_period: float = 8., - -): - r"""Judge spiking state with a squarewave fourier series. - - The forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - - Backward function: - - .. math:: - - g'(x) = \sum_{i=1}^n\frac{4\cos\left((2 * i - 1.) * 2\pi * x / T\right)}{T} - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> bp.visualize.get_figure(1, 1, 4, 6) - >>> xs = bm.linspace(-3, 3, 1000) - >>> for n in [2, 4, 8]: - >>> f = bm.surrogate.SquarewaveFourierSeries(n=n) - >>> grads1 = bm.vector_grad(f)(xs) - >>> plt.plot(bm.as_numpy(xs), bm.as_numpy(grads1), label=f'n={n}') - >>> plt.legend() - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - n: int - t_period: float - - - Returns:: - - out: jax.Array - The spiking state. - - """ - - return SquarewaveFourierSeries(n=n, t_period=t_period)(x) - - -class S2NN(Surrogate): - """Judge spiking state with the S2NN surrogate spiking function. - - See Also:: - - s2nn - """ - - def __init__(self, alpha=4., beta=1., epsilon=1e-8): - super().__init__() - self.alpha = alpha - self.beta = beta - self.epsilon = epsilon - - def surrogate_fun(self, x): - x = as_jax(x) - z = jnp.where(x < 0., - sci.special.expit(x * self.alpha), - self.beta * jnp.log(jnp.abs((x + 1.)) + self.epsilon) + 0.5) - return z - - def surrogate_grad(self, x): - x = as_jax(x) - sg = sci.special.expit(self.alpha * x) - dx = jnp.where(x < 0., self.alpha * sg * (1. - sg), self.beta / (x + 1.)) - return dx - - def __repr__(self): - return f'{self.__class__.__name__}(alpha={self.alpha}, beta={self.beta}, epsilon={self.epsilon})' - - -def s2nn( - x: Union[jax.Array, Array], - alpha: float = 4., - beta: float = 1., - epsilon: float = 1e-8, - -): - r"""Judge spiking state with the S2NN surrogate spiking function [1]_. - - The forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - - Backward function: - - .. math:: - - \begin{split}g'(x) = \begin{cases} - \alpha * (1 - \mathrm{sigmoid} (\alpha x)) \mathrm{sigmoid} (\alpha x), x < 0 \\ - \frac{\beta}{(x + 1)}, x \ge 0 - \end{cases}\end{split} - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> bp.visualize.get_figure(1, 1, 4, 6) - >>> xs = bm.linspace(-3, 3, 1000) - >>> grads = bm.vector_grad(bm.surrogate.s2nn)(xs, 4., 1.) - >>> plt.plot(bm.as_numpy(xs), bm.as_numpy(grads), label=r'$\alpha=4, \beta=1$') - >>> grads = bm.vector_grad(bm.surrogate.s2nn)(xs, 8., 2.) - >>> plt.plot(bm.as_numpy(xs), bm.as_numpy(grads), label=r'$\alpha=8, \beta=2$') - >>> plt.legend() - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - alpha: float - The param that controls the gradient when ``x < 0``. - beta: float - The param that controls the gradient when ``x >= 0`` - epsilon: float - Avoid nan - - - Returns:: - - out: jax.Array - The spiking state. - - References:: - - .. [1] Suetake, Kazuma et al. “S2NN: Time Step Reduction of Spiking Surrogate Gradients for Training Energy Efficient Single-Step Neural Networks.” ArXiv abs/2201.10879 (2022): n. pag. - - """ - return S2NN(alpha=alpha, beta=beta, epsilon=epsilon)(x) - - -class QPseudoSpike(Surrogate): - """Judge spiking state with the q-PseudoSpike surrogate function. - - See Also:: - - q_pseudo_spike - """ - - def __init__(self, alpha=2.): - super().__init__() - self.alpha = alpha - - def surrogate_grad(self, x): - x = as_jax(x) - dx = jnp.power(1 + 2 / (self.alpha - 1) * jnp.abs(x), -self.alpha) - return dx - - def surrogate_fun(self, x): - x = as_jax(x) - z = jnp.where(x < 0., - 0.5 * jnp.power(1 - 2 / (self.alpha - 1) * x, 1 - self.alpha), - 1. - 0.5 * jnp.power(1 + 2 / (self.alpha - 1) * jnp.abs(x), 1 - self.alpha)) - return z - - def __repr__(self): - return f'{self.__class__.__name__}(alpha={self.alpha})' - - -def q_pseudo_spike( - x: Union[jax.Array, Array], - alpha: float = 2., - -): - r"""Judge spiking state with the q-PseudoSpike surrogate function [1]_. - - The forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - - Backward function: - - .. math:: - - g'(x) = (1+\frac{2|x|}{\alpha-1})^{-\alpha} - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> xs = bm.linspace(-3, 3, 1000) - >>> bp.visualize.get_figure(1, 1, 4, 6) - >>> for alpha in [0.5, 1., 2., 4.]: - >>> grads = bm.vector_grad(bm.surrogate.q_pseudo_spike)(xs, alpha) - >>> plt.plot(bm.as_numpy(xs), bm.as_numpy(grads), label=r'$\alpha=$' + str(alpha)) - >>> plt.legend() - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - alpha: float - The parameter to control tail fatness of gradient. - - - Returns:: - - out: jax.Array - The spiking state. - - References:: - - .. [1] Herranz-Celotti, Luca and Jean Rouat. “Surrogate Gradients Design.” ArXiv abs/2202.00282 (2022): n. pag. - """ - return QPseudoSpike(alpha=alpha)(x) - - -class LeakyRelu(Surrogate): - """Judge spiking state with the Leaky ReLU function. - - See Also:: - - leaky_relu - """ - - def __init__(self, alpha=0.1, beta=1.): - super().__init__() - self.alpha = alpha - self.beta = beta - - def surrogate_fun(self, x): - x = as_jax(x) - return jnp.where(x < 0., self.alpha * x, self.beta * x) - - def surrogate_grad(self, x): - x = as_jax(x) - dx = jnp.where(x < 0., self.alpha, self.beta) - return dx - - def __repr__(self): - return f'{self.__class__.__name__}(alpha={self.alpha}, beta={self.beta})' - - -def leaky_relu( - x: Union[jax.Array, Array], - alpha: float = 0.1, - beta: float = 1., - -): - r"""Judge spiking state with the Leaky ReLU function. - - The forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - - Backward function: - - .. math:: - - \begin{split}g'(x) = - \begin{cases} - \beta, & x \geq 0 \\ - \alpha, & x < 0 \\ - \end{cases}\end{split} - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> xs = bm.linspace(-3, 3, 1000) - >>> bp.visualize.get_figure(1, 1, 4, 6) - >>> grads = bm.vector_grad(bm.surrogate.leaky_relu)(xs, 0., 1.) - >>> plt.plot(bm.as_numpy(xs), bm.as_numpy(grads), label=r'$\alpha=0., \beta=1.$') - >>> plt.legend() - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - alpha: float - The parameter to control the gradient when :math:`x < 0`. - beta: float - The parameter to control the gradient when :math:`x >= 0`. - - - Returns:: - - out: jax.Array - The spiking state. - """ - return LeakyRelu(alpha=alpha, beta=beta)(x) - - -class LogTailedRelu(Surrogate): - """Judge spiking state with the Log-tailed ReLU function. - - See Also:: - - log_tailed_relu - """ - - def __init__(self, alpha=0.): - super().__init__() - self.alpha = alpha - - def surrogate_fun(self, x): - x = as_jax(x) - z = jnp.where(x > 1, - jnp.log(x), - jnp.where(x > 0, - x, - self.alpha * x)) - return z - - def surrogate_grad(self, x): - x = as_jax(x) - dx = jnp.where(x > 1, - 1 / x, - jnp.where(x > 0, - 1., - self.alpha)) - return dx - - def __repr__(self): - return f'{self.__class__.__name__}(alpha={self.alpha})' - - -def log_tailed_relu( - x: Union[jax.Array, Array], - alpha: float = 0., - -): - r"""Judge spiking state with the Log-tailed ReLU function [1]_. - - The forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - - Backward function: - - .. math:: - - \begin{split}g'(x) = - \begin{cases} - \alpha, & x \leq 0 \\ - 1, & 0 < x \leq 0 \\ - \frac{1}{x}, x > 1 \\ - \end{cases}\end{split} - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> xs = bm.linspace(-3, 3, 1000) - >>> bp.visualize.get_figure(1, 1, 4, 6) - >>> grads = bm.vector_grad(bm.surrogate.leaky_relu)(xs, 0., 1.) - >>> plt.plot(bm.as_numpy(xs), bm.as_numpy(grads), label=r'$\alpha=0., \beta=1.$') - >>> plt.legend() - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - alpha: float - The parameter to control the gradient. - - - Returns:: - - out: jax.Array - The spiking state. - - References:: - - .. [1] Cai, Zhaowei et al. “Deep Learning with Low Precision by Half-Wave Gaussian Quantization.” 2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR) (2017): 5406-5414. - """ - return LogTailedRelu(alpha=alpha)(x) - - -class ReluGrad(Surrogate): - """Judge spiking state with the ReLU gradient function. - - See Also:: - - relu_grad - """ - - def __init__(self, alpha=0.3, width=1.): - super().__init__() - self.alpha = alpha - self.width = width - - def surrogate_grad(self, x): - x = as_jax(x) - dx = jnp.maximum(self.alpha * self.width - jnp.abs(x) * self.alpha, 0) - return dx - - def __repr__(self): - return f'{self.__class__.__name__}(alpha={self.alpha}, width={self.width})' - - -def relu_grad( - x: Union[jax.Array, Array], - alpha: float = 0.3, - width: float = 1., -): - r"""Spike function with the ReLU gradient function [1]_. - - The forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - Backward function: - - .. math:: - - g'(x) = \text{ReLU}(\alpha * (\mathrm{width}-|x|)) - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> xs = bm.linspace(-3, 3, 1000) - >>> bp.visualize.get_figure(1, 1, 4, 6) - >>> for s in [0.5, 1.]: - >>> for w in [1, 2.]: - >>> grads = bm.vector_grad(bm.surrogate.relu_grad)(xs, s, w) - >>> plt.plot(bm.as_numpy(xs), bm.as_numpy(grads), label=r'$\alpha=$' + f'{s}, width={w}') - >>> plt.legend() - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - alpha: float - The parameter to control the gradient. - width: float - The parameter to control the width of the gradient. - - Returns:: - - out: jax.Array - The spiking state. - - References:: - - .. [1] Neftci, E. O., Mostafa, H. & Zenke, F. Surrogate gradient learning in spiking neural networks. IEEE Signal Process. Mag. 36, 61–63 (2019). - """ - return ReluGrad(alpha=alpha, width=width)(x) - - -class GaussianGrad(Surrogate): - """Judge spiking state with the Gaussian gradient function. - - See Also:: - - gaussian_grad - """ - - def __init__(self, sigma=0.5, alpha=0.5): - super().__init__() - self.sigma = sigma - self.alpha = alpha - - def surrogate_grad(self, x): - x = as_jax(x) - dx = jnp.exp(-(x ** 2) / (2 * jnp.power(self.sigma, 2))) / (jnp.sqrt(2 * jnp.pi) * self.sigma) - return self.alpha * dx - - def __repr__(self): - return f'{self.__class__.__name__}(alpha={self.alpha}, sigma={self.sigma})' - - -def gaussian_grad( - x: Union[jax.Array, Array], - sigma: float = 0.5, - alpha: float = 0.5, -): - r"""Spike function with the Gaussian gradient function [1]_. - - The forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - Backward function: - - .. math:: - - g'(x) = \alpha * \text{gaussian}(x, 0., \sigma) - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> xs = bm.linspace(-3, 3, 1000) - >>> bp.visualize.get_figure(1, 1, 4, 6) - >>> for s in [0.5, 1., 2.]: - >>> grads = bm.vector_grad(bm.surrogate.gaussian_grad)(xs, s, 0.5) - >>> plt.plot(bm.as_numpy(xs), bm.as_numpy(grads), label=r'$\alpha=0.5, \sigma=$' + str(s)) - >>> plt.legend() - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - sigma: float - The parameter to control the variance of gaussian distribution. - alpha: float - The parameter to control the scale of the gradient. - - Returns:: - - out: jax.Array - The spiking state. - - References:: - - .. [1] Yin, B., Corradi, F. & Bohté, S.M. Accurate and efficient time-domain classification with adaptive spiking recurrent neural networks. Nat Mach Intell 3, 905–913 (2021). - """ - return GaussianGrad(sigma=sigma, alpha=alpha)(x) - - -class MultiGaussianGrad(Surrogate): - """Judge spiking state with the multi-Gaussian gradient function. - - See Also:: - - multi_gaussian_grad - """ - - def __init__(self, h=0.15, s=6.0, sigma=0.5, scale=0.5): - super().__init__() - self.h = h - self.s = s - self.sigma = sigma - self.scale = scale - - def surrogate_grad(self, x): - x = as_jax(x) - g1 = jnp.exp(-x ** 2 / (2 * jnp.power(self.sigma, 2))) / (jnp.sqrt(2 * jnp.pi) * self.sigma) - g2 = jnp.exp(-(x - self.sigma) ** 2 / (2 * jnp.power(self.s * self.sigma, 2)) - ) / (jnp.sqrt(2 * jnp.pi) * self.s * self.sigma) - g3 = jnp.exp(-(x + self.sigma) ** 2 / (2 * jnp.power(self.s * self.sigma, 2)) - ) / (jnp.sqrt(2 * jnp.pi) * self.s * self.sigma) - dx = g1 * (1. + self.h) - g2 * self.h - g3 * self.h - return self.scale * dx - - def __repr__(self): - return f'{self.__class__.__name__}(h={self.h}, s={self.s}, sigma={self.sigma}, scale={self.scale})' - - -def multi_gaussian_grad( - x: Union[jax.Array, Array], - h: float = 0.15, - s: float = 6.0, - sigma: float = 0.5, - scale: float = 0.5, -): - r"""Spike function with the multi-Gaussian gradient function [1]_. - - The forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - Backward function: - - .. math:: - - \begin{array}{l} - g'(x)=(1+h){{{\mathcal{N}}}}(x, 0, {\sigma }^{2}) - -h{{{\mathcal{N}}}}(x, \sigma,{(s\sigma )}^{2})- - h{{{\mathcal{N}}}}(x, -\sigma ,{(s\sigma )}^{2}) - \end{array} - - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> xs = bm.linspace(-3, 3, 1000) - >>> bp.visualize.get_figure(1, 1, 4, 6) - >>> grads = bm.vector_grad(bm.surrogate.multi_gaussian_grad)(xs) - >>> plt.plot(bm.as_numpy(xs), bm.as_numpy(grads)) - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - h: float - The hyper-parameters of approximate function - s: float - The hyper-parameters of approximate function - sigma: float - The gaussian sigma. - scale: float - The gradient scale. - - Returns:: - - out: jax.Array - The spiking state. - - References:: - - .. [1] Yin, B., Corradi, F. & Bohté, S.M. Accurate and efficient time-domain classification with adaptive spiking recurrent neural networks. Nat Mach Intell 3, 905–913 (2021). - """ - return MultiGaussianGrad(h=h, s=s, sigma=sigma, scale=scale)(x) - - -class InvSquareGrad(Surrogate): - """Judge spiking state with the inverse-square surrogate gradient function. - - See Also:: - - inv_square_grad - """ - - def __init__(self, alpha=100.): - super().__init__() - self.alpha = alpha - - def surrogate_grad(self, x): - dx = 1. / (self.alpha * jnp.abs(x) + 1.0) ** 2 - return dx - - def __repr__(self): - return f'{self.__class__.__name__}(alpha={self.alpha})' - - -def inv_square_grad( - x: Union[jax.Array, Array], - alpha: float = 100. -): - r"""Spike function with the inverse-square surrogate gradient. - - Forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - Backward function: - - .. math:: - - g'(x) = \frac{1}{(\alpha * |x| + 1.) ^ 2} - - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> xs = bm.linspace(-1, 1, 1000) - >>> for alpha in [1., 10., 100.]: - >>> grads = bm.vector_grad(bm.surrogate.inv_square_grad)(xs, alpha) - >>> plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha)) - >>> plt.legend() - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - alpha: float - Parameter to control smoothness of gradient - - Returns:: - - out: jax.Array - The spiking state. - """ - return InvSquareGrad(alpha=alpha)(x) - - -class SlayerGrad(Surrogate): - """Judge spiking state with the slayer surrogate gradient function. - - See Also:: - - slayer_grad - """ - - def __init__(self, alpha=1.): - super().__init__() - self.alpha = alpha - - def surrogate_grad(self, x): - dx = jnp.exp(-self.alpha * jnp.abs(x)) - return dx - - def __repr__(self): - return f'{self.__class__.__name__}(alpha={self.alpha})' - - -def slayer_grad( - x: Union[jax.Array, Array], - alpha: float = 1. -): - r"""Spike function with the slayer surrogate gradient function. - - Forward function: - - .. math:: - - g(x) = \begin{cases} - 1, & x \geq 0 \\ - 0, & x < 0 \\ - \end{cases} - - Backward function: - - .. math:: - - g'(x) = \exp(-\alpha |x|) - - - .. plot:: - :include-source: True - - >>> import brainpy as bp - >>> import brainpy.math as bm - >>> import matplotlib.pyplot as plt - >>> bp.visualize.get_figure(1, 1, 4, 6) - >>> xs = bm.linspace(-3, 3, 1000) - >>> for alpha in [0.5, 1., 2., 4.]: - >>> grads = bm.vector_grad(bm.surrogate.slayer_grad)(xs, alpha) - >>> plt.plot(xs, grads, label=r'$\alpha$=' + str(alpha)) - >>> plt.legend() - >>> plt.show() - - Parameters:: - - x: jax.Array, Array - The input data. - alpha: float - Parameter to control smoothness of gradient - - Returns:: - - out: jax.Array - The spiking state. - - References:: - - .. [1] Shrestha, S. B. & Orchard, G. Slayer: spike layer error reassignment in time. In Advances in Neural Information Processing Systems Vol. 31, 1412–1421 (NeurIPS, 2018). - """ - return SlayerGrad(alpha=alpha)(x) diff --git a/brainpy/math/surrogate/_two_inputs.py b/brainpy/math/surrogate/_two_inputs.py deleted file mode 100644 index d4bebff59..000000000 --- a/brainpy/math/surrogate/_two_inputs.py +++ /dev/null @@ -1,67 +0,0 @@ -# -*- coding: utf-8 -*- -# Copyright 2025 BrainX 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 Union - -import jax -import jax.numpy as jnp - -from brainpy.math.interoperability import as_jax -from brainpy.math.ndarray import Array -from ._utils import vjp_custom - -__all__ = [ - 'inv_square_grad2', - 'relu_grad2', -] - - -@vjp_custom(['x_new', 'x_old'], dict(alpha=100.)) -def inv_square_grad2( - x_new: Union[jax.Array, Array], - x_old: Union[jax.Array, Array], - alpha: float -): - x_new_comp = x_new >= 0 - x_old_comp = x_old < 0 - z = jnp.asarray(jnp.logical_and(x_new_comp, x_old_comp), dtype=x_new.dtype) - - def grad(dz): - dz = as_jax(dz) - dx_new = (dz / (alpha * jnp.abs(x_new) + 1.0) ** 2) * jnp.asarray(x_old_comp, dtype=x_old.dtype) - dx_old = -(dz / (alpha * jnp.abs(x_old) + 1.0) ** 2) * jnp.asarray(x_new_comp, dtype=x_new.dtype) - return dx_new, dx_old, None - - return z, grad - - -@vjp_custom(['x_new', 'x_old'], dict(alpha=.3, width=1.)) -def relu_grad2( - x_new: Union[jax.Array, Array], - x_old: Union[jax.Array, Array], - alpha: float, - width: float, -): - x_new_comp = x_new >= 0 - x_old_comp = x_old < 0 - z = jnp.asarray(jnp.logical_and(x_new_comp, x_old_comp), dtype=x_new.dtype) - - def grad(dz): - dz = as_jax(dz) - dx_new = (dz * jnp.maximum(width - jnp.abs(x_new), 0) * alpha) * jnp.asarray(x_old_comp, dtype=x_old.dtype) - dx_old = -(dz * jnp.maximum(width - jnp.abs(x_old), 0) * alpha) * jnp.asarray(x_new_comp, dtype=x_new.dtype) - return dx_new, dx_old, None, None - - return z, grad diff --git a/brainpy/math/surrogate/_utils.py b/brainpy/math/surrogate/_utils.py deleted file mode 100644 index 1fdbaad75..000000000 --- a/brainpy/math/surrogate/_utils.py +++ /dev/null @@ -1,146 +0,0 @@ -# -*- coding: utf-8 -*- -# Copyright 2025 BrainX 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. -# ============================================================================== -import inspect -import itertools -from functools import partial -from typing import Dict, Callable, Sequence - -import jax - -from brainpy import check -from brainpy._errors import UnsupportedError -from brainpy.math.ndarray import Array as Array - -__all__ = [ - 'get_default', - 'make_return', - 'vjp_custom', -] - - -def get_default(x, default): - if x is None: - return default, False - else: - return x, True - - -def make_return(r, *args): - if isinstance(r, (tuple, list)): - r = tuple(r) - else: - r = [r] - for a in args: - if a: - r += [None] - return tuple(r) - - -def _get_args(f): - reduced_args = [] - for name, par in inspect.signature(f).parameters.items(): - if par.kind is inspect.Parameter.POSITIONAL_OR_KEYWORD: - reduced_args.append(par.name) - - elif par.kind is inspect.Parameter.VAR_POSITIONAL: - reduced_args.append(f'*{par.name}') - - elif par.kind is inspect.Parameter.KEYWORD_ONLY: - raise UnsupportedError() - elif par.kind is inspect.Parameter.POSITIONAL_ONLY: - raise UnsupportedError() - elif par.kind is inspect.Parameter.VAR_KEYWORD: # TODO - raise UnsupportedError() - else: - raise UnsupportedError() - return reduced_args - - -class VJPCustom(object): - def __init__(self, - func: Callable, - args: Sequence[str], - defaults: Dict = None, - statics: Dict = None, ): - if statics is None: statics = dict() - if defaults is None: defaults = dict() - assert isinstance(statics, dict) - assert isinstance(defaults, dict) - assert callable(func) - check.is_sequence(args, elem_type=str) - - self.n_args = len(args) - self.func = func - self.args = args - self.defaults = tuple(defaults.items()) - self.statics = statics - self.all_args = _get_args(func) - - for k in statics: - if k not in defaults: - raise KeyError(f'{k} defined as "static_args" should provide its default value in "defaults"') - self._cached = {} - if len(statics): - static_vals = list(statics.values()) - products = list(itertools.product(*static_vals, repeat=1)) - for args in products: - string = self._str_static_arg(dict(zip(self.statics.keys(), args))) - self._cached[string] = jax.custom_gradient(partial(self.func, - **dict(zip(self.statics.keys(), args)))) - else: - self._cached[''] = jax.custom_gradient(self.func) - - def _str_static_arg(self, args: Dict): - r = [] - for k in self.statics: - r.append(f'{k}={args[k]}') - return '-'.join(r) - - def __call__(self, *args, **kwargs): - args = list(args) - kwargs = dict(kwargs) - for k in self.args[len(args):]: - if k not in kwargs: - raise ValueError(f'Must provide {k} for function {self.func}') - args.append(kwargs.pop(k)) - for k, v in self.defaults[len(args) - self.n_args:]: - if k not in kwargs: - args.append(v) - else: - args.append(kwargs.pop(k)) - if len(kwargs): - raise KeyError(f'Unknown arguments {kwargs} for function {self.func}') - dynamics = [] - statics = dict() - for k, v in zip(self.all_args, args): - if isinstance(v, Array): v = v.value - if k in self.statics: - statics[k] = v - else: - dynamics.append(v) - return self._cached[self._str_static_arg(statics)](*dynamics) - - -def vjp_custom(args: Sequence[str], defaults: Dict, statics: Dict = None): - """Generalize a customized gradient function as a general Python function. - """ - - def wrapper(fun): - obj = VJPCustom(fun, args, defaults, statics) - obj.__doc__ = fun.__doc__ - return obj - - return wrapper diff --git a/brainpy/math/surrogate/one_input_test.py b/brainpy/math/surrogate/one_input_test.py deleted file mode 100644 index dede3bcb2..000000000 --- a/brainpy/math/surrogate/one_input_test.py +++ /dev/null @@ -1,64 +0,0 @@ -# -*- coding: utf-8 -*- -# Copyright 2025 BrainX 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. -# ============================================================================== -import jax -from absl.testing import parameterized - -import brainpy.math as bm -from brainpy.math.surrogate import _one_input as one_input - - -class TestOneInputGrad(parameterized.TestCase): - def __init__(self, *args, platform='cpu', **kwargs): - super(TestOneInputGrad, self).__init__(*args, **kwargs) - bm.set_platform(platform) - print() - - @parameterized.named_parameters( - dict(testcase_name=f'{name}_x64={x64}', - func=getattr(one_input, name), - x64=x64) - for name in one_input.__all__ - for x64 in [True, False] - ) - def test_bm_grad(self, func, x64): - if x64: - bm.enable_x64() - - xs = bm.arange(-3, 3, 0.005) - grads = bm.vector_grad(func)(xs) - self.assertTrue(grads.size == xs.size) - - if x64: - bm.disable_x64() - - @parameterized.named_parameters( - dict(testcase_name=f'{name}_x64={x64}', - func=getattr(one_input, name), - x64=x64, ) - for name in one_input.__all__ - for x64 in [True, False] - ) - def test_jax_vjp(self, func, x64): - if x64: - bm.enable_x64() - - xs = bm.arange(-3, 3, 0.005) - primals, f_vjp = jax.vjp(func, xs) - grad2 = f_vjp(jax.numpy.ones_like(xs)) - self.assertTrue(grad2[0].size == xs.size) - - if x64: - bm.disable_x64() diff --git a/brainpy/math/surrogate/two_inputs_test.py b/brainpy/math/surrogate/two_inputs_test.py deleted file mode 100644 index 1471c8b55..000000000 --- a/brainpy/math/surrogate/two_inputs_test.py +++ /dev/null @@ -1,64 +0,0 @@ -# -*- coding: utf-8 -*- -# Copyright 2025 BrainX 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. -# ============================================================================== -import jax -from absl.testing import parameterized - -import brainpy.math as bm -from brainpy.math.surrogate import _two_inputs as two_inputs - - -class TestTwoInputsGrad(parameterized.TestCase): - def __init__(self, *args, platform='cpu', **kwargs): - super(TestTwoInputsGrad, self).__init__(*args, **kwargs) - bm.set_platform(platform) - print() - - @parameterized.named_parameters( - dict(testcase_name=f'{name}_x64={x64}', - func=getattr(two_inputs, name), - x64=x64) - for name in two_inputs.__all__ - for x64 in [True, False] - ) - def test_bm_grad(self, func, x64): - if x64: - bm.enable_x64() - - xs = bm.arange(-3, 3, 0.005) - grads = bm.vector_grad(func)(xs[:-1], xs[1:]) - self.assertTrue(grads.size == xs.size - 1) - - if x64: - bm.disable_x64() - - @parameterized.named_parameters( - dict(testcase_name=f'{name}_x64={x64}', - func=getattr(two_inputs, name), - x64=x64, ) - for name in two_inputs.__all__ - for x64 in [True, False] - ) - def test_jax_vjp(self, func, x64): - if x64: - bm.enable_x64() - - xs = bm.arange(-3, 3, 0.005) - primals, f_vjp = jax.vjp(func, xs[:-1], xs[1:]) - grad2 = f_vjp(jax.numpy.ones(xs.size - 1)) - self.assertTrue(grad2[0].size == xs.size - 1) - - if x64: - bm.disable_x64() diff --git a/docs/apis/brainpy.math.surrogate.rst b/docs/apis/brainpy.math.surrogate.rst index e4c7d2b91..f32120e2d 100644 --- a/docs/apis/brainpy.math.surrogate.rst +++ b/docs/apis/brainpy.math.surrogate.rst @@ -1,8 +1,13 @@ ``brainpy.math.surrogate``: Surrogate Gradient Functions ================================================================= -.. currentmodule:: brainpy.math.surrogate -.. automodule:: brainpy.math.surrogate +.. note:: + + ``brainpy.math.surrogate`` reuses :mod:`braintools.surrogate`: + ``brainpy.math.surrogate`` is an alias of ``braintools.surrogate``. + +.. currentmodule:: braintools.surrogate +.. automodule:: braintools.surrogate .. autosummary:: :toctree: generated/ @@ -44,5 +49,3 @@ multi_gaussian_grad SlayerGrad slayer_grad - inv_square_grad2 - relu_grad2 \ No newline at end of file diff --git a/docs/auto_generater.py b/docs/auto_generater.py index 2f076f7cb..e3884aa25 100644 --- a/docs/auto_generater.py +++ b/docs/auto_generater.py @@ -579,7 +579,7 @@ def generate_math_docs(): 'sparse': ('``brainpy.math.sparse`` module: Sparse Operators', 'brainpy.math.sparse'), 'event': ('``brainpy.math.event`` module: Event-driven Operators', 'brainpy.math.event'), 'jitconn': ('``brainpy.math.jitconn`` module: Just-In-Time Connectivity Operators', 'brainpy.math.jitconn'), - 'surrogate': ('``brainpy.math.surrogate`` module: Surrogate Gradient Functions', 'brainpy.math.surrogate'), + 'surrogate': ('``brainpy.math.surrogate`` module: Surrogate Gradient Functions', 'braintools.surrogate'), 'random': ('``brainpy.math.random`` module: Random Number Generations', 'brainpy.math.random'), 'linalg': ('``brainpy.math.linalg`` module: Linear algebra', 'brainpy.math.linalg'), 'fft': ('``brainpy.math.fft`` module: Discrete Fourier Transform', 'brainpy.math.fft'), diff --git a/pyproject.toml b/pyproject.toml index 0edf82774..69a289d17 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,7 +40,7 @@ dependencies = [ "brainstate>=0.5.1", "brainunit>=0.2.0", "brainevent>=0.0.7", - "braintools>=0.0.9", + "braintools>=0.2.0", 'brainpy_state>=0.0.3', ] diff --git a/requirements.txt b/requirements.txt index e02aa230d..23904ac62 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,7 @@ numpy>=1.15 brainunit brainevent>=0.0.7 -braintools>=0.0.9 +braintools>=0.2.0 brainstate>=0.5.1 brainpy_state>=0.0.3 jax