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

brainstate/nn/event/csr.py

@@ -1,312 +0,0 @@
-# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-from typing import Union, Callable, Optional
-
-import brainunit as u
-import jax
-import jax.numpy as jnp
-import numpy as np
-
-from brainstate._state import ParamState, State
-from brainstate.init import param
-from brainstate.mixin import Mode, Training
-from brainstate.nn._base import DnnLayer
-from brainstate.typing import ArrayLike
-from ._misc import IntScalar
-
-__all__ = [
-  'EventCSR',
-]
-
-
-class EventCSR(DnnLayer):
-  """
-  The EventCSR module implements a fixed probability connection with CSR sparse data structure.
-
-  Parameters
-  ----------
-  n_pre : int
-      Number of pre-synaptic neurons.
-  n_post : int
-      Number of post-synaptic neurons.
-  weight : float or callable or jax.Array or brainunit.Quantity
-      Maximum synaptic conductance.
-  name : str, optional
-      Name of the module.
-  mode : brainstate.mixin.Mode, optional
-      Mode of the module.
-  """
-
-  def __init__(
-      self,
-      n_pre: IntScalar,
-      n_post: IntScalar,
-      indptr: ArrayLike,
-      indices: ArrayLike,
-      weight: Union[Callable, ArrayLike],
-      name: Optional[str] = None,
-      mode: Optional[Mode] = None,
-      grad_mode: str = 'vjp'
-  ):
-    super().__init__(name=name, mode=mode)
-
-    self.in_size = n_pre
-    self.out_size = n_post
-    self.n_pre = n_pre
-    self.n_post = n_post
-
-    # gradient mode
-    assert grad_mode in ['vjp', 'jvp'], f"Unsupported grad_mode: {grad_mode}"
-    self.grad_mode = grad_mode
-
-    # CSR data structure
-    indptr = jnp.asarray(indptr)
-    indices = jnp.asarray(indices)
-    assert indptr.ndim == 1, f"indptr must be 1D. Got: {indptr.ndim}"
-    assert indices.ndim == 1, f"indices must be 1D. Got: {indices.ndim}"
-    assert indptr.size == n_pre + 1, f"indptr must have size {n_pre + 1}. Got: {indptr.size}"
-    self.indptr = indptr
-    self.indices = indices
-
-    # maximum synaptic conductance
-    weight = param(weight, (len(indices),), allow_none=False)
-    # if callable(weight):
-    #   pass
-    # else:
-    #   if u.math.size(weight) != 1 and u.math.size(weight) != len(self.indices):
-    #     raise ValueError(f"weight must be 1D or 2D with size {len(self.indices)}. Got: {u.math.size(weight)}")
-    if self.mode.has(Training):
-      weight = ParamState(weight)
-    self.weight = weight
-
-  def update(self, spk: jax.Array) -> Union[jax.Array, u.Quantity]:
-    weight = self.weight.value if isinstance(self.weight, State) else self.weight
-    if len(self.indices) == 0:
-      r = u.math.zeros(spk.shape[:-1] + (self.n_post,),
-                       dtype=weight.dtype,
-                       unit=u.get_unit(weight) * u.get_unit(spk))
-      return u.maybe_decimal(r)
-
-    device_kind = jax.devices()[0].platform  # spk.device.device_kind
-    if device_kind == 'cpu':
-      return cpu_event_csr(
-        u.math.asarray(spk),
-        u.math.asarray(self.indptr),
-        u.math.asarray(self.indices),
-        u.math.asarray(weight),
-        n_post=self.n_post, grad_mode=self.grad_mode
-      )
-    elif device_kind in ['gpu', 'tpu']:
-      raise NotImplementedError()
-    else:
-      raise ValueError(f"Unsupported device: {device_kind}")
-
-
-def cpu_event_csr(
-    spk: jax.Array,
-    indptr: jax.Array,
-    indices: jax.Array,
-    weight: Union[u.Quantity, jax.Array],
-    *,
-    n_post: int,
-    grad_mode: str = 'vjp'
-) -> Union[u.Quantity, jax.Array]:
-  """
-  The EventCSR module implements a fixed probability connection with CSR sparse data structure.
-
-  Parameters
-  ----------
-  spk : jax.Array
-      Spike events.
-  indptr : jax.Array
-      Index pointer of post connected neurons.
-  indices : jax.Array
-      Indices of post connected neurons.
-  weight : brainunit.Quantity or jax.Array
-      Maximum synaptic conductance.
-  n_post : int
-      Number of post-synaptic neurons.
-  grad_mode : str, optional
-      Gradient mode. Default is 'vjp'. Can be 'vjp' or 'jvp'.
-
-  Returns
-  -------
-  post_data : brainunit.Quantity or jax.Array
-      Post synaptic data.
-  """
-  unit = u.get_unit(weight)
-  weight = u.get_mantissa(weight)
-
-  def mv(spk_vector):
-    assert spk_vector.ndim == 1, f"spk must be 1D. Got: {spk.ndim}"
-    if grad_mode == 'vjp':
-      post_data = _cpu_event_csr_mv_vjp(spk_vector, indptr, indices, weight, n_post)
-    elif grad_mode == 'jvp':
-      post_data = _cpu_event_csr_mv_jvp(spk_vector, indptr, indices, weight, n_post)
-    else:
-      raise ValueError(f"Unsupported grad_mode: {grad_mode}")
-    return post_data
-
-  assert spk.ndim >= 1, f"spk must be at least 1D. Got: {spk.ndim}"
-  assert weight.ndim in [1, 0], f"g_max must be 1D or 0D. Got: {weight.ndim}"
-  assert indices.ndim == 1, f"indices must be 1D. Got: {indices.ndim}"
-
-  if spk.ndim == 1:
-    post_data = mv(spk)
-  else:
-    shape = spk.shape[:-1]
-    post_data = jax.vmap(mv)(u.math.reshape(spk, (-1, spk.shape[-1])))
-    post_data = u.math.reshape(post_data, shape + post_data.shape[-1:])
-  return u.maybe_decimal(u.Quantity(post_data, unit=unit))
-
-
-# --------------
-# Implementation
-# --------------
-
-
-def _cpu_event_csr_mv(
-    spk: jax.Array,
-    indptr: jax.Array,
-    indices: jax.Array,
-    weight: Union[u.Quantity, jax.Array],
-    n_post: int
-) -> jax.Array:
-  bool_x = spk.dtype == jnp.bool_
-  homo_w = jnp.size(weight) == 1
-
-  def add_fn(post_val, i_start, i_end, sp):
-    def body_fn(x):
-      post, i = x
-      i_post = indices[i]
-      w = weight if homo_w else weight[i]
-      w = w if bool_x else w * sp
-      post = post.at[i_post].add(w)
-      return post, i + 1
-
-    return jax.lax.while_loop(lambda x: x[1] < i_end, body_fn, (post_val, i_start))[0]
-
-  def scan_fn(post, i):
-    sp = spk[i]  # pre-synaptic spike event
-    if bool_x:
-      post = jax.lax.cond(sp, lambda: add_fn(post, indptr[i], indptr[i + 1], sp), lambda: post)
-    else:
-      post = jax.lax.cond(sp == 0., lambda: post, lambda: add_fn(post, indptr[i], indptr[i + 1], sp))
-    return post, None
-
-  return jax.lax.scan(scan_fn, jnp.zeros((n_post,), dtype=weight.dtype), np.arange(len(spk)))[0]
-
-
-# --------------
-# VJP
-# --------------
-
-def _cpu_event_csr_mv_fwd(
-    spk: jax.Array,
-    indptr: jax.Array,
-    indices: jax.Array,
-    weight: Union[u.Quantity, jax.Array],
-    n_post: int
-):
-  return _cpu_event_csr_mv(spk, indptr, indices, weight, n_post=n_post), (spk, weight)
-
-
-def _cpu_event_csr_mv_bwd(indptr, indices, n_post, res, ct):
-  spk, weight = res
-  homo = jnp.size(weight) == 1
-  bool_spk = spk.dtype == jnp.bool_
-
-  # ∂L/∂spk = ∂L/∂y * ∂y/∂spk
-  def fn_spk(i_pre):
-    def body_fn(x):
-      r, i = x
-      i_post = indices[i]
-      r = r + (ct[i_post] if homo else ct[i_post] * weight[i])
-      return r, i + 1
-
-    p = jax.lax.while_loop(lambda x: x[1] < indptr[i_pre + 1], body_fn, (0., indptr[i_pre]))[0]
-    p = p * weight if homo else p
-    return p
-
-  ct_spk = jax.vmap(fn_spk)(np.arange(len(spk)))
-
-  # ∂L/∂w = ∂L/∂y * ∂y/∂w
-  if homo:  # scalar
-    ct_gmax = _cpu_event_csr_mv(spk, indptr, indices, jnp.asarray(1.), n_post=n_post)
-    ct_gmax = jnp.inner(ct, ct_gmax)
-  else:
-    def single_post(dw, i_pre):
-      def body_fn(x):
-        dw, i = x
-        i_post = indices[i]
-        dw = dw.at[i].add(ct[i_post] if bool_spk else ct[i_post] * spk[i_pre])
-        return dw, i + 1
-
-      return jax.lax.while_loop(lambda x: x[1] < indptr[i_pre + 1], body_fn, (dw, indptr[i_pre]))[0]
-
-    def fn_w(dw, i_pre):
-      sp = spk[i_pre]
-      if bool_spk:
-        return jax.lax.cond(sp, lambda: single_post(dw, i_pre), lambda: dw), None
-      else:
-        return jax.lax.cond(sp == 0., lambda: dw, lambda: single_post(dw, i_pre)), None
-
-    ct_gmax = jax.lax.scan(fn_w, jnp.zeros_like(weight), np.arange(len(spk)))[0]
-  return ct_spk, ct_gmax
-
-
-_cpu_event_csr_mv_vjp = jax.custom_vjp(_cpu_event_csr_mv, nondiff_argnums=(1, 2, 4))
-_cpu_event_csr_mv_vjp.defvjp(_cpu_event_csr_mv_fwd, _cpu_event_csr_mv_bwd)
-
-
-# --------------
-# JVP
-# --------------
-
-
-def _cpu_event_csr_mv_jvp_rule(indptr, indices, n_post, primals, tangents):
-  # forward pass
-  spk, weight = primals
-  y = _cpu_event_csr_mv(spk, indptr, indices, weight, n_post=n_post)
-
-  # forward gradients
-  spk_dot, weight_dot = tangents
-  homo_w = jnp.size(weight) == 1
-
-  # ∂y/∂gmax
-  dweight = _cpu_event_csr_mv(spk, indptr, indices, weight_dot, n_post=n_post)
-
-  # ∂y/∂gspk
-  def scan_fn(post, i_pre):
-    def while_fn(x):
-      p, i, sp = x
-      i_post = indices[i]
-      p = p.at[i_post].add(sp if homo_w else sp * weight[i])
-      return p, i + 1, sp
-
-    post = jax.lax.while_loop(lambda x: x[1] < indptr[i_pre + 1],
-                              while_fn,
-                              (post, indptr[i_pre], spk_dot[i_pre]))[0]
-
-    return post, None
-
-  dspk = jax.lax.scan(scan_fn, jnp.zeros((n_post,), dtype=weight.dtype), np.arange(len(spk)))[0]
-  dspk = (dspk * weight) if homo_w else dspk
-  return y, dweight + dspk
-
-
-_cpu_event_csr_mv_jvp = jax.custom_jvp(_cpu_event_csr_mv, nondiff_argnums=(1, 2, 4))
-_cpu_event_csr_mv_jvp.defjvp(_cpu_event_csr_mv_jvp_rule)

brainstate/nn/event/csr_test.py

@@ -1,118 +0,0 @@
-# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-
-import jax.numpy
-import jax.numpy as jnp
-import numpy as np
-from absl.testing import parameterized
-
-import brainstate as bst
-from brainstate.nn.event.csr import EventCSR
-
-
-def _get_csr(n_pre, n_post, prob):
-  n_conn = int(n_post * prob)
-  indptr = np.arange(n_pre + 1) * n_conn
-  indices = np.random.randint(0, n_post, (n_pre * n_conn,))
-  return indptr, indices
-
-
-def true_fn(x, w, indices, indptr, n_out):
-  homo_w = jnp.size(w) == 1
-
-  post = jnp.zeros((n_out,))
-  for i_pre in range(x.shape[0]):
-    ids = indices[indptr[i_pre]: indptr[i_pre + 1]]
-    post = post.at[ids].add(w * x[i_pre] if homo_w else w[indptr[i_pre]: indptr[i_pre + 1]] * x[i_pre])
-  return post
-
-
-class TestFixedProbCSR(parameterized.TestCase):
-  @parameterized.product(
-    homo_w=[True, False],
-  )
-  def test1(self, homo_w):
-    x = bst.random.rand(20) < 0.1
-    indptr, indices = _get_csr(20, 40, 0.1)
-    m = EventCSR(20, 40, indptr, indices, 1.5 if homo_w else bst.init.Normal())
-    y = m(x)
-    y2 = true_fn(x, m.weight, indices, indptr, 40)
-    self.assertTrue(jnp.allclose(y, y2))
-
-  @parameterized.product(
-    bool_x=[True, False],
-    homo_w=[True, False]
-  )
-  def test_vjp(self, bool_x, homo_w):
-    n_in = 20
-    n_out = 30
-    if bool_x:
-      x = jax.numpy.asarray(bst.random.rand(n_in) < 0.3, dtype=float)
-    else:
-      x = bst.random.rand(n_in)
-
-    indptr, indices = _get_csr(n_in, n_out, 0.1)
-    fn = bst.nn.EventCSR(n_in, n_out, indptr, indices, 1.5 if homo_w else bst.init.Normal())
-    w = fn.weight
-
-    def f(x, w):
-      fn.weight = w
-      return fn(x).sum()
-
-    r = jax.grad(f, argnums=(0, 1))(x, w)
-
-    # -------------------
-    # TRUE gradients
-
-    def f2(x, w):
-      return true_fn(x, w, indices, indptr, n_out).sum()
-
-    r2 = jax.grad(f2, argnums=(0, 1))(x, w)
-    self.assertTrue(jnp.allclose(r[0], r2[0]))
-    self.assertTrue(jnp.allclose(r[1], r2[1]))
-
-  @parameterized.product(
-    bool_x=[True, False],
-    homo_w=[True, False]
-  )
-  def test_jvp(self, bool_x, homo_w):
-    n_in = 20
-    n_out = 30
-    if bool_x:
-      x = jax.numpy.asarray(bst.random.rand(n_in) < 0.3, dtype=float)
-    else:
-      x = bst.random.rand(n_in)
-
-    indptr, indices = _get_csr(n_in, n_out, 0.1)
-    fn = EventCSR(n_in, n_out, indptr, indices,
-                  1.5 if homo_w else bst.init.Normal(), grad_mode='jvp')
-    w = fn.weight
-
-    def f(x, w):
-      fn.weight = w
-      return fn(x)
-
-    o1, r1 = jax.jvp(f, (x, w), (jnp.ones_like(x), jnp.ones_like(w)))
-
-    # -------------------
-    # TRUE gradients
-
-    def f2(x, w):
-      return true_fn(x, w, indices, indptr, n_out)
-
-    o2, r2 = jax.jvp(f2, (x, w), (jnp.ones_like(x), jnp.ones_like(w)))
-    self.assertTrue(jnp.allclose(r1, r2))
-    self.assertTrue(jnp.allclose(o1, o2))

brainstate/nn/event/fixed_probability.py

@@ -1,276 +0,0 @@
-# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-
-from typing import Union, Callable, Optional
-
-import brainunit as u
-import jax
-import jax.numpy as jnp
-import numpy as np
-
-from brainstate._state import ParamState, State
-from brainstate.init import param
-from brainstate.mixin import Mode, Training
-from brainstate.nn._base import DnnLayer
-from brainstate.random import RandomState
-from brainstate.transform import for_loop
-from brainstate.typing import ArrayLike
-from ._misc import FloatScalar, IntScalar
-
-__all__ = [
-  'EventFixedProb',
-]
-
-
-class EventFixedProb(DnnLayer):
-  """
-  The EventFixedProb module implements a fixed probability connection with CSR sparse data structure.
-
-  Parameters
-  ----------
-  n_pre : int
-      Number of pre-synaptic neurons.
-  n_post : int
-      Number of post-synaptic neurons.
-  prob : float
-      Probability of connection.
-  weight : float or callable or jax.Array or brainunit.Quantity
-      Maximum synaptic conductance.
-  allow_multi_conn : bool, optional
-      Whether multiple connections are allowed from a single pre-synaptic neuron.
-      Default is True, meaning that a value of ``a`` can be selected multiple times.
-  prob : float
-      Probability of connection.
-  name : str, optional
-      Name of the module.
-  mode : brainstate.mixin.Mode, optional
-      Mode of the module.
-  """
-
-  def __init__(
-      self,
-      n_pre: IntScalar,
-      n_post: IntScalar,
-      prob: FloatScalar,
-      weight: Union[Callable, ArrayLike],
-      allow_multi_conn: bool = True,
-      seed: Optional[int] = None,
-      name: Optional[str] = None,
-      mode: Optional[Mode] = None,
-      grad_mode: str = 'vjp'
-  ):
-    super().__init__(name=name, mode=mode)
-    self.n_pre = n_pre
-    self.n_post = n_post
-    self.in_size = n_pre
-    self.out_size = n_post
-
-    self.n_conn = int(n_post * prob)
-    if self.n_conn < 1:
-      raise ValueError(f"The number of connections must be at least 1. Got: int({n_post} * {prob}) = {self.n_conn}")
-
-    assert grad_mode in ['vjp', 'jvp'], f"Unsupported grad_mode: {grad_mode}"
-    self.grad_mode = grad_mode
-
-    # indices of post connected neurons
-    if allow_multi_conn:
-      self.indices = np.random.RandomState(seed).randint(0, n_post, size=(self.n_pre, self.n_conn))
-    else:
-      rng = RandomState(seed)
-      self.indices = for_loop(lambda i: rng.choice(n_post, size=(self.n_conn,), replace=False), np.arange(n_pre))
-
-    # maximum synaptic conductance
-    weight = param(weight, (self.n_pre, self.n_conn), allow_none=False)
-    if self.mode.has(Training):
-      weight = ParamState(weight)
-    self.weight = weight
-
-  def update(self, spk: jax.Array) -> Union[jax.Array, u.Quantity]:
-    weight = self.weight.value if isinstance(self.weight, State) else self.weight
-    device_kind = jax.devices()[0].platform  # spk.device.device_kind
-    if device_kind == 'cpu':
-      return cpu_event_fixed_prob(u.math.asarray(self.indices),
-                                  u.math.asarray(weight),
-                                  u.math.asarray(spk),
-                                  n_post=self.n_post, grad_mode=self.grad_mode)
-    elif device_kind in ['gpu', 'tpu']:
-      raise NotImplementedError()
-    else:
-      raise ValueError(f"Unsupported device: {device_kind}")
-
-
-def cpu_event_fixed_prob(
-    indices: jax.Array,
-    weight: Union[u.Quantity, jax.Array],
-    spk: jax.Array,
-    *,
-    n_post: int,
-    grad_mode: str = 'vjp'
-) -> Union[u.Quantity, jax.Array]:
-  """
-  The EventFixedProb module implements a fixed probability connection with CSR sparse data structure.
-
-  Parameters
-  ----------
-  n_post : int
-      Number of post-synaptic neurons.
-  weight : brainunit.Quantity or jax.Array
-      Maximum synaptic conductance.
-  spk : jax.Array
-      Spike events.
-  indices : jax.Array
-      Indices of post connected neurons.
-  grad_mode : str, optional
-      Gradient mode. Default is 'vjp'. Can be 'vjp' or 'jvp'.
-
-  Returns
-  -------
-  post_data : brainunit.Quantity or jax.Array
-      Post synaptic data.
-  """
-  unit = u.get_unit(weight)
-  weight = u.get_mantissa(weight)
-  indices = jnp.asarray(indices)
-  spk = jnp.asarray(spk)
-
-  def mv(spk_vector):
-    assert spk_vector.ndim == 1, f"spk must be 1D. Got: {spk.ndim}"
-    if grad_mode == 'vjp':
-      post_data = _cpu_event_fixed_prob_mv_vjp(indices, weight, spk_vector, n_post)
-    elif grad_mode == 'jvp':
-      post_data = _cpu_event_fixed_prob_mv_jvp(indices, weight, spk_vector, n_post)
-    else:
-      raise ValueError(f"Unsupported grad_mode: {grad_mode}")
-    return post_data
-
-  assert spk.ndim >= 1, f"spk must be at least 1D. Got: {spk.ndim}"
-  assert weight.ndim in [2, 0], f"weight must be 2D or 0D. Got: {weight.ndim}"
-  assert indices.ndim == 2, f"indices must be 2D. Got: {indices.ndim}"
-
-  if spk.ndim == 1:
-    post_data = mv(spk)
-  else:
-    shape = spk.shape[:-1]
-    post_data = jax.vmap(mv)(u.math.reshape(spk, (-1, spk.shape[-1])))
-    post_data = u.math.reshape(post_data, shape + post_data.shape[-1:])
-  return u.maybe_decimal(u.Quantity(post_data, unit=unit))
-
-
-# -------------------
-# CPU Implementation
-# -------------------
-
-
-def _cpu_event_fixed_prob_mv(indices, g_max, spk, n_post: int) -> jax.Array:
-  def scan_fn(post, i):
-    w = g_max if jnp.size(g_max) == 1 else g_max[i]
-    ids = indices[i]
-    sp = spk[i]
-    if spk.dtype == jnp.bool_:
-      post = jax.lax.cond(sp, lambda: post.at[ids].add(w), lambda: post)
-    else:
-      post = jax.lax.cond(sp == 0., lambda: post, lambda: post.at[ids].add(w * sp))
-    return post, None
-
-  return jax.lax.scan(scan_fn, jnp.zeros((n_post,), dtype=g_max.dtype), np.arange(len(spk)))[0]
-
-
-# --------------
-# VJP
-# --------------
-
-def _cpu_event_fixed_prob_mv_fwd(indices, g_max, spk, n_post):
-  return _cpu_event_fixed_prob_mv(indices, g_max, spk, n_post=n_post), (g_max, spk)
-
-
-def _cpu_event_fixed_prob_mv_bwd(indices, n_post, res, ct):
-  weight, spk = res
-
-  # ∂L/∂spk = ∂L/∂y * ∂y/∂spk
-  homo = jnp.size(weight) == 1
-  if homo:  # homogeneous weight
-    ct_spk = jax.vmap(lambda idx: jnp.sum(ct[idx] * weight))(indices)
-  else:  # heterogeneous weight
-    ct_spk = jax.vmap(lambda idx, w: jnp.inner(ct[idx], w))(indices, weight)
-
-  # ∂L/∂w = ∂L/∂y * ∂y/∂w
-  if homo:  # scalar
-    ct_gmax = _cpu_event_fixed_prob_mv(indices, jnp.asarray(1.), spk, n_post=n_post)
-    ct_gmax = jnp.inner(ct, ct_gmax)
-  else:
-    def scan_fn(d_gmax, i):
-      if spk.dtype == jnp.bool_:
-        d_gmax = jax.lax.cond(spk[i], lambda: d_gmax.at[i].add(ct[indices[i]]), lambda: d_gmax)
-      else:
-        d_gmax = jax.lax.cond(spk[i] == 0., lambda: d_gmax, lambda: d_gmax.at[i].add(ct[indices[i]] * spk[i]))
-      return d_gmax, None
-
-    ct_gmax = jax.lax.scan(scan_fn, jnp.zeros_like(weight), np.arange(len(spk)))[0]
-  return ct_gmax, ct_spk
-
-
-_cpu_event_fixed_prob_mv_vjp = jax.custom_vjp(_cpu_event_fixed_prob_mv, nondiff_argnums=(0, 3))
-_cpu_event_fixed_prob_mv_vjp.defvjp(_cpu_event_fixed_prob_mv_fwd, _cpu_event_fixed_prob_mv_bwd)
-
-
-# --------------
-# JVP
-# --------------
-
-
-def _cpu_event_fixed_prob_mv_jvp_rule(indices, n_post, primals, tangents):
-  # forward pass
-  weight, spk = primals
-  y = _cpu_event_fixed_prob_mv(indices, weight, spk, n_post=n_post)
-
-  # forward gradients
-  gmax_dot, spk_dot = tangents
-
-  # ∂y/∂gmax
-  dgmax = _cpu_event_fixed_prob_mv(indices, gmax_dot, spk, n_post=n_post)
-
-  def scan_fn(post, i):
-    ids = indices[i]
-    w = weight if jnp.size(weight) == 1 else weight[i]
-    post = post.at[ids].add(w * spk_dot[i])
-    return post, None
-
-  # ∂y/∂gspk
-  dspk = jax.lax.scan(scan_fn, jnp.zeros((n_post,), dtype=weight.dtype), np.arange(len(spk)))[0]
-  return y, dgmax + dspk
-
-
-_cpu_event_fixed_prob_mv_jvp = jax.custom_jvp(_cpu_event_fixed_prob_mv, nondiff_argnums=(0, 3))
-_cpu_event_fixed_prob_mv_jvp.defjvp(_cpu_event_fixed_prob_mv_jvp_rule)
-
-
-
-
-
-
-def _gpu_event_fixed_prob_mv(indices, g_max, spk, n_post: int) -> jax.Array:
-  def scan_fn(post, i):
-    w = g_max if jnp.size(g_max) == 1 else g_max[i]
-    ids = indices[i]
-    sp = spk[i]
-    if spk.dtype == jnp.bool_:
-      post = jax.lax.cond(sp, lambda: post.at[ids].add(w), lambda: post)
-    else:
-      post = jax.lax.cond(sp == 0., lambda: post, lambda: post.at[ids].add(w * sp))
-    return post, None
-
-  return jax.lax.scan(scan_fn, jnp.zeros((n_post,), dtype=g_max.dtype), np.arange(len(spk)))[0]
-
-