Just-in-Time Connection Matrices#
Just-in-time (JIT) connectivity represents a reproducible random matrix through distribution parameters and a seed. Connections are generated during an operation instead of being stored as an explicit edge list.
import brainevent
import jax
import jax.numpy as jnp
Why Generate Connections Just in Time?#
An explicit sparse matrix stores every realized edge. A JIT matrix instead stores a compact generative specification. This can reduce persistent connectivity storage and make seeded experiments reproducible, but it trades storage for on-demand generation. Runtime and temporary-memory behavior depend on the kernel, shape, probability, dtype, backend, and operation.
Homogeneous-Weight JIT Connectivity#
JITCScalarR gives every realized edge one scalar weight. The tuple is (weight, connection_probability, seed). The same specification generates the same matrix.
scalar_r = brainevent.JITCScalarR((0.2, 0.25, 42), shape=(8, 5))
scalar_r_again = brainevent.JITCScalarR((0.2, 0.25, 42), shape=(8, 5))
assert jnp.array_equal(scalar_r.todense(), scalar_r_again.todense())
print(scalar_r.todense())
---------------------------------------------------------------------------
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/_jit_scalar/float.py:404, in _jitc_homo_matrix_numba_kernel(corder, transpose, **kwargs)
380 """
381 Build a Numba CPU kernel for generating a dense JIT scalar connectivity matrix.
382
(...) 402 A kernel function with signature ``(weight, clen, seed) -> tuple``.
403 """
--> 404 import numba
405 _rng = get_numba_light_rng_funcs()
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/_jit_scalar/float.py:137, in jits()
136 clen = _initialize_conn_length(prob)
--> 137 res = jits_p_call(
138 weight,
139 clen,
140 seed,
141 shape=shape,
142 transpose=transpose,
143 corder=corder,
144 backend=backend,
145 )[0]
146 return u.maybe_decimal(res * unitd)
File /opt/hostedtoolcache/Python/3.14.7/x64/lib/python3.14/site-packages/brainevent/_jit_scalar/float.py:733, in jits_p_call()
727 out_info = (
728 jax.ShapeDtypeStruct(shape[::-1], dtype=weight.dtype)
729 if transpose else
730 jax.ShapeDtypeStruct(shape, dtype=weight.dtype)
731 )
--> 733 return jits_p(
734 weight,
735 clen,
736 seed,
737 outs=[out_info],
738 weight_info=jax.ShapeDtypeStruct(weight.shape, weight.dtype),
739 clen_info=jax.ShapeDtypeStruct(clen.shape, clen.dtype),
740 seed_info=jax.ShapeDtypeStruct(seed.shape, seed.dtype),
741 out_info=out_info,
742 shape=shape,
743 transpose=transpose,
744 corder=corder,
745 backend=backend,
746 )
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 'float_jitc_homo_matrix' on platform 'cpu': ModuleNotFoundError: No module named 'numba'. Available backend(s) for platform 'cpu': ['numba']. Switch with backend='<name>' on this call, or float_jitc_homo_matrix.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[2], line 3
1 scalar_r = brainevent.JITCScalarR((0.2, 0.25, 42), shape=(8, 5))
2 scalar_r_again = brainevent.JITCScalarR((0.2, 0.25, 42), shape=(8, 5))
----> 3 assert jnp.array_equal(scalar_r.todense(), scalar_r_again.todense())
4 print(scalar_r.todense())
File /opt/hostedtoolcache/Python/3.14.7/x64/lib/python3.14/site-packages/brainevent/_jit_scalar/main.py:651, in JITCScalarR.todense(self)
605 def todense(self) -> Union[jax.Array, u.Quantity]:
606 """
607 Convert the sparse scalar-weight matrix to dense format.
608
(...) 649 >>> dense_matrix.shape # (10, 4)
650 """
--> 651 return jits(
652 self.weight, self.prob, self.seed,
653 shape=self.shape, transpose=False, corder=self.corder,
654 backend=self.backend,
655 )
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 'float_jitc_homo_matrix' on platform 'cpu': ModuleNotFoundError: No module named 'numba'. Available backend(s) for platform 'cpu': ['numba']. Switch with backend='<name>' on this call, or float_jitc_homo_matrix.set_default('cpu', '<name>') / brainevent.set_backend('cpu', '<name>') to change the default.
Normally Distributed JIT Connectivity#
JITCNormalR draws realized weights from a normal distribution. Its tuple is (location, scale, connection_probability, seed). Distribution parameters specify the generator; sample statistics of a small realized matrix need not equal them exactly.
normal_r = brainevent.JITCNormalR((0.0, 0.1, 0.25, 43), shape=(8, 5))
print(normal_r.todense())
Uniformly Distributed JIT Connectivity#
JITCUniformR draws realized weights uniformly between lower and upper bounds. Its tuple is (low, high, connection_probability, seed).
uniform_r = brainevent.JITCUniformR((-0.1, 0.3, 0.25, 44), shape=(8, 5))
print(uniform_r.todense())
Memory and Performance Trade-offs#
The persistent JIT representation contains parameters and a seed rather than arrays for every realized edge. That observation does not establish end-to-end memory savings: generated intermediates, compilation caches, and backend-specific kernels also contribute. Likewise, a stored CSR matrix may be faster when the same explicit edges are reused. Measure peak memory, compilation time, and synchronized steady-state execution separately on the target hardware.
Build a Large Random Network#
The example keeps the shape moderate enough for documentation CI while illustrating a layer whose connectivity is generated from a seed.
n_pre, n_post = 512, 128
large_random = brainevent.JITCNormalR((0.0, 0.05, 0.02, 2024), shape=(n_pre, n_post))
events = brainevent.BinaryArray(jnp.arange(n_pre) % 29 == 0)
forward = jax.jit(lambda x: x @ large_random)
postsynaptic_input = forward(events)
postsynaptic_input.block_until_ready()
print(postsynaptic_input.shape)
Row- and Column-Oriented Connectivity#
The R and C variants expose complementary orientations. Transposition preserves the generated connectivity while swapping the matrix axes. Choose the orientation that matches the dominant multiplication direction, then verify with the actual workload.
scalar_c = scalar_r.transpose()
small_events = brainevent.BinaryArray(jnp.array([1, 0, 1, 0, 1, 0, 0, 1], dtype=bool))
row_result = small_events @ scalar_r
column_result = scalar_c @ small_events
assert jnp.allclose(row_result, column_result)
print(type(scalar_r).__name__, type(scalar_c).__name__)
Summary and Next Steps#
JIT matrices define reproducible generated connectivity with scalar, normal, or uniform weights. They avoid storing an explicit persistent edge list, but their full performance and memory costs are workload-dependent. Compare them with CSR and CSC and Fixed Connection Count Structures using the representation that matches the scientific constraint.