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()
../_images/80c71d1fe6cf754c5314db77b0f02cb0996aa862757624e147827740033ac9f9.png

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.