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())
[[0.2 0. 0. 0.2 0. ]
[0. 0. 0. 0.2 0. ]
[0. 0.2 0. 0.2 0.2]
[0. 0. 0. 0. 0. ]
[0. 0.2 0. 0.2 0.2]
[0. 0. 0. 0. 0. ]
[0. 0.2 0. 0.2 0. ]
[0. 0. 0. 0.2 0. ]]
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())
[[ 0.00650178 0. 0.07174362 0. -0.03097547]
[ 0.00697076 0. 0.02893138 0. 0. ]
[ 0. 0. 0. 0. 0.16859469]
[ 0. 0. 0. -0.09673285 0. ]
[ 0.0821434 0. 0. 0. 0. ]
[ 0. 0. 0. 0. 0. ]
[ 0. 0. 0. 0. 0. ]
[ 0.0060613 0. 0. 0. -0.25178513]]
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())
[[0. 0. 0.2630237 0.02689505 0. ]
[0. 0. 0. 0. 0. ]
[0. 0. 0. 0. 0. ]
[0. 0. 0. 0. 0. ]
[0. 0. 0. 0. 0. ]
[0.02787633 0. 0.18712984 0. 0. ]
[0.02428703 0. 0. 0. 0.01985388]
[0. 0.13362844 0. 0.07484462 0. ]]
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)
(128,)
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__)
JITCScalarR JITCScalarC
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.