Quickstart#
BrainEvent connects two ideas: Data describes how neural connectivity is stored or generated, while Events describes sparse, discrete activity and the operations driven by it. This notebook takes you from those concepts to a visible spike pattern and a first event-driven matrix multiplication. For setup instructions, see Installation.
What BrainEvent Computes#
A BinaryArray wraps boolean or 0/1 activity. When it participates in matrix multiplication, BrainEvent processes the active entries as events while preserving the numerical result of ordinary dense multiplication.
Why Event-Driven Computation?#
Event-driven kernels can avoid work associated with inactive entries. The benefit depends on event density, matrix shape, backend, hardware, compilation state, and memory behavior; sparsity alone does not guarantee a speedup.
Import BrainEvent#
import brainevent
import brainstate
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
print(f"BrainEvent {brainevent.__version__}")
print(f"JAX backend: {jax.default_backend()}")
BrainEvent 0.2.1
JAX backend: cpu
Create and Visualize Binary Events#
Create 80 time steps for 20 neurons. Each active entry represents a spike.
brainstate.random.seed(7)
spike_train = brainstate.random.bernoulli(0.12, size=(80, 20))
events = brainevent.BinaryArray(spike_train)
print("event shape:", events.shape)
print("total active events:", int(spike_train.sum()))
event shape: (80, 20)
total active events: 175
time_index, neuron_index = jnp.nonzero(spike_train)
fig, ax = plt.subplots(figsize=(8, 3))
ax.scatter(time_index, neuron_index, s=8)
ax.set(xlabel="time step", ylabel="neuron", title="Binary spike events")
plt.tight_layout()
plt.show()
Run Your First Event-Driven Matrix Multiplication#
Multiply the event batch by dense connectivity weights. The ordinary JAX product provides a correctness reference.
weights = jnp.linspace(-0.5, 0.5, 60, dtype=jnp.float32).reshape(20, 3)
event_output = events @ weights
dense_output = spike_train @ weights
print("output shape:", event_output.shape)
print("matches dense result:", bool(jnp.allclose(event_output, dense_output)))
---------------------------------------------------------------------------
ModuleNotFoundError Traceback (most recent call last)
File /opt/hostedtoolcache/Python/3.14.7/x64/lib/python3.14/site-packages/brainevent/_op/main.py:569, in XLACustomKernel._register_fallback_lowering.<locals>.fallback_kernel_fn(*args, **kwargs)
568 try:
--> 569 kernel = entry.kernel_generator(**kwargs)
570 except Exception as exc:
File /opt/hostedtoolcache/Python/3.14.7/x64/lib/python3.14/site-packages/brainevent/_dense/binary.py:585, in _binary_densemm_numba_kernel(spk_info, weight_info, transpose, **kwargs)
579 def _binary_densemm_numba_kernel(
580 spk_info: jax.ShapeDtypeStruct,
581 weight_info: jax.ShapeDtypeStruct,
582 transpose: bool,
583 **kwargs
584 ):
--> 585 import numba
587 if transpose:
588 # weights[k,m].T @ spikes[k,n] -> out[m,n]
589 # primitive args: (weights, spikes)
ModuleNotFoundError: No module named 'numba'
The above exception was the direct cause of the following exception:
JaxStackTraceBeforeTransformation Traceback (most recent call last)
File /opt/hostedtoolcache/Python/3.14.7/x64/lib/python3.14/site-packages/brainevent/_dense/binary.py:575, in binary_densemm()
574 spk_val = spk_val.astype(jnp.bool_)
--> 575 r = binary_densemm_p_call(weight_val, spk_val, transpose=transpose, backend=backend)
576 return u.maybe_decimal(r[0] * wunit * spkunit)
File /opt/hostedtoolcache/Python/3.14.7/x64/lib/python3.14/site-packages/brainevent/_dense/binary.py:904, in binary_densemm_p_call()
903 out = jax.ShapeDtypeStruct([weights.shape[0], spikes.shape[1]], weights.dtype)
--> 904 return binary_densemm_p(
905 weights,
906 spikes,
907 outs=[out],
908 transpose=transpose,
909 spk_info=jax.ShapeDtypeStruct(spikes.shape, spikes.dtype),
910 weight_info=jax.ShapeDtypeStruct(weights.shape, weights.dtype),
911 backend=backend,
912 )
File /opt/hostedtoolcache/Python/3.14.7/x64/lib/python3.14/site-packages/brainevent/_op/main.py:338, in __call__()
337 flat_outs, tree_def = abstract_arguments(outs)
--> 338 r = self.primitive.bind(*ins, **kwargs, outs=tuple(flat_outs))
339 if len(r) != len(flat_outs):
JaxStackTraceBeforeTransformation: brainevent.KernelCompilationError: Backend 'numba' failed to construct a kernel for primitive 'binary_densemm' on platform 'cpu': ModuleNotFoundError: No module named 'numba'. Available backend(s) for platform 'cpu': ['numba', 'jax_raw']. Switch with backend='<name>' on this call, or binary_densemm.set_default('cpu', '<name>') / brainevent.set_backend('cpu', '<name>') to change the default.
The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.
--------------------
The above exception was the direct cause of the following exception:
KernelCompilationError Traceback (most recent call last)
Cell In[4], line 2
1 weights = jnp.linspace(-0.5, 0.5, 60, dtype=jnp.float32).reshape(20, 3)
----> 2 event_output = events @ weights
3 dense_output = spike_train @ weights
4
5 print("output shape:", event_output.shape)
File /opt/hostedtoolcache/Python/3.14.7/x64/lib/python3.14/site-packages/brainevent/_event/binary.py:212, in BinaryArray.__matmul__(self, oc)
208 return binary_densemv(oc, self.value, transpose=True)
209 else: # self.ndim == 2
210 # self[m,k] @ oc[k,n]: use weights=oc[k,n], spikes=self.value.T[k,m]
211 # gives oc.T @ self.value.T = [n,m], then .T = [m,n]
--> 212 return binary_densemm(oc, self.value.T, transpose=True).T
213 else:
214 return oc.__rmatmul__(self)
File /opt/hostedtoolcache/Python/3.14.7/x64/lib/python3.14/site-packages/brainevent/_misc.py:1770, in NameScope.__call__(self, *args, **kwargs)
1768 backend = kwargs.pop('backend', None)
1769 jit_fn = self._get_jit_fn(backend)
-> 1770 return jit_fn(*args, **kwargs)
[... skipping hidden 18 frame]
File /opt/hostedtoolcache/Python/3.14.7/x64/lib/python3.14/site-packages/brainevent/_op/main.py:571, in XLACustomKernel._register_fallback_lowering.<locals>.fallback_kernel_fn(*args, **kwargs)
569 kernel = entry.kernel_generator(**kwargs)
570 except Exception as exc:
--> 571 raise KernelCompilationError(
572 f"Backend '{backend_to_use}' failed to construct a kernel "
573 f"for primitive '{self.name}' on platform '{platform}': "
574 f"{type(exc).__name__}: {exc}." + _construction_error_suffix()
575 ) from exc
577 try:
578 return kernel(*args)
KernelCompilationError: Backend 'numba' failed to construct a kernel for primitive 'binary_densemm' on platform 'cpu': ModuleNotFoundError: No module named 'numba'. Available backend(s) for platform 'cpu': ['numba', 'jax_raw']. Switch with backend='<name>' on this call, or binary_densemm.set_default('cpu', '<name>') / brainevent.set_backend('cpu', '<name>') to change the default.
Use BrainEvent with JAX Transformations#
BrainEvent arrays compose with JAX transformations. The compiled function below performs the same event-driven multiplication.
@jax.jit
def apply_events(binary_events, matrix):
return binary_events @ matrix
compiled_output = apply_events(events, weights)
print("compiled result matches:", bool(jnp.allclose(compiled_output, dense_output)))
Summary and Next Steps#
You created binary spike events, visualized their activity, and verified an event-driven matrix multiplication against dense JAX computation. Continue with Data for CSR/CSC, fixed-count, and just-in-time connectivity; Events for event representations and event-triggered updates; or Custom operators to extend BrainEvent with new kernels.