Process supplied histories

Choose history processing or rollout

Use scan_histories when the complete native input history is available. The input tensor has shape (neurons, time, input_channels); one neuron is the same interface with neurons=1. With a supported official CUDA Mamba model, the selective state-space recurrence is evaluated by the fused temporal scan rather than a Python loop over native samples.

The supplied history must contain the inputs themselves, in the channel order, sign convention and sampling cadence used to train the model. A connected network’s future recurrent events still depend on spikes generated during simulation, so they must be produced by causal population rollout. A recorded network history can subsequently be processed through this interface, but scanning it does not generate a new closed-loop network trajectory.

Your inputs	Execution path	Temporal state
Complete known histories	scan_histories	Fresh sequence state for each call
Samples arriving over time for fixed streams	Single-neuron streaming	Caller retains streaming state between samples/chunks
Inputs generated by recurrent population activity	Connected population	Runtime retains learned state, clock and delay queue

The last tensor dimension is the model’s native input-channel width, model.num_input. It is not the number of retained recurrent edges in a connected graph. Prepare inputs in the checkpoint’s channel representation before using this interface, or use the Lite supplied-history fitting interface for its explicit contact-to-channel mapping.

Run one neuron

After installation, download history_scan_example.py into your working directory. It supplies a small untrained AxoPyTorchMamba and deterministic signed input histories for CPU checks; this portable fallback is sequential and provides no fused-scan timing or biological-accuracy evidence.

python history_scan_example.py

pytorch-sequential-reference False
single (1, 12, 2)
population (3, 12, 2)
direct outputs, untouched neuron rows and input/model gradients agree

To process one neuron’s history, preserve the leading neuron axis and provide its declared morphology index. The following runnable example also checks that the wrapper returns the same prediction as the underlying model on that history:

import torch
from axosim import scan_histories
from history_scan_example import build_reference_model, make_histories

model = build_reference_model()
inputs = make_histories(neurons=1)           # (1, 12, 8)
morphology_indices = torch.tensor([0])
with torch.no_grad():
    result = scan_histories(
        model, inputs, morphology_indices=morphology_indices,
        require_parallel=False,
    )
    direct = model(inputs, morphology_indices=morphology_indices)
torch.testing.assert_close(result.prediction, direct, atol=1e-6, rtol=1e-5)
print(result.prediction.shape, result.backend.name, result.backend.parallel)

torch.Size([1, 12, 2]) pytorch-sequential-reference False

Inspect result.backend rather than inferring temporal parallelism from the Mamba name or tensor shape. The default require_parallel=True rejects unsupported execution instead of silently accepting the sequential fallback. For application predictions, replace the example model with trained weights and preserve their actual morphology vocabulary and input preprocessing.

Process a population sharing weights

Each neuron row receives its own complete input history and morphology assignment while sharing the provided model weights. neuron_batch_size bounds the number of rows processed together; it divides the neuron axis only, so the entire time axis and its within-history context remain intact.

import torch
from axosim import scan_histories
from history_scan_example import build_reference_model, make_histories

model = build_reference_model()
inputs = make_histories(neurons=3)           # (3, 12, 8)
morphology_indices = torch.tensor([0, 1, 0])
with torch.no_grad():
    result = scan_histories(
        model, inputs, morphology_indices=morphology_indices,
        neuron_batch_size=2, require_parallel=False,
    )
    direct = model(inputs, morphology_indices=morphology_indices)
    torch.testing.assert_close(result.prediction, direct, atol=1e-6, rtol=1e-5)

    changed = inputs.clone()
    changed[0, :, 0] += 0.5
    altered = scan_histories(
        model, changed, morphology_indices=morphology_indices,
        require_parallel=False,
    ).prediction
torch.testing.assert_close(altered[1:], result.prediction[1:], atol=1e-6, rtol=1e-5)
assert not torch.allclose(altered[0], result.prediction[0])
print(result.prediction.shape)

torch.Size([3, 12, 2])

Changing one supplied row leaves other neurons’ predictions unchanged because this call does not route new spikes between rows. Their recorded inputs may already reflect connected activity; the supplied values, rather than a live connectome, define this computation. Match each index to model.config.morphology_ids when the model uses morphology conditioning. For strict CUDA FP32 comparisons across neuron batch sizes, keep TF32 settings consistent and use numerical tolerances rather than bitwise equality.

Require the fused selective scan

Load a trusted official Mamba1 checkpoint on CUDA and keep the default parallel requirement. This example uses your local checkpoint path; official Mamba dependencies are described in installation. Predictions use the model’s spike-logit and soma-target coordinates.

import torch
from axosim import scan_histories
from axosim.checkpoint import load_checkpoint

model, metadata = load_checkpoint("path/to/trained-axomamba.pt", map_location="cuda")
model = model.eval()
device = next(model.parameters()).device
dtype = next(model.parameters()).dtype
inputs = torch.zeros(1, 8_192, model.num_input, device=device, dtype=dtype)
morphology_indices = (
    torch.zeros(1, dtype=torch.long, device=device)
    if model.config.morphology_ids else None
)
with torch.no_grad():
    result = scan_histories(
        model, inputs, morphology_indices=morphology_indices,
        require_parallel=True,
    )
assert result.backend.parallel
print(result.prediction.shape, result.backend.name)

The reported backend is official-mamba1-fused when the call succeeds. The model must contain the supported official Mamba1 blocks with use_fast_path=True, and both fused selective-scan and causal-convolution kernels must be available. This zero-input call verifies execution/device placement, not accuracy; submit measured or simulated native histories for evaluation.

The CPU portable Mamba backend is a separate reference implementation with a different checkpoint format, so do not load official released weights into the small fallback used above. Setting require_parallel=False permits that reference model but never substitutes a backend for an official checkpoint. The scan reference specifies the exact validation and failure conditions.

Preserve gradients and sequence boundaries

scan_histories retains autograd and the caller’s training/evaluation mode. For fitting or input-sensitivity analysis, leave gradient tracking enabled. This runnable CPU check verifies gradients through both supplied histories and the model readout:

import torch
from axosim import scan_histories
from history_scan_example import build_reference_model, make_histories

model = build_reference_model()
inputs = make_histories(neurons=3).requires_grad_()
result = scan_histories(
    model, inputs, morphology_indices=torch.tensor([0, 1, 0]),
    neuron_batch_size=2, require_parallel=False,
)
result.prediction.square().mean().backward()
assert inputs.grad is not None and torch.isfinite(inputs.grad).all()
assert model.readout.weight.grad is not None
assert torch.isfinite(model.readout.weight.grad).all()
print(inputs.grad.shape, model.readout.weight.grad.shape)

torch.Size([3, 12, 8]) torch.Size([2, 4])

Each call starts a fresh history state and returns predictions, rather than a retained streaming state. Splitting the time axis into separate calls therefore does not continue the preceding history. Reuse the existing streaming state when later samples continue the same neuron streams. Gradients here describe a model conditioned on supplied inputs; the interface does not add gradients through discrete connected-event generation.

Next steps

Use AxoBench evaluation to measure accuracy under the benchmark’s input/target protocol, or training to optimize model weights. The scan API reference lists exact arguments, outputs, backend metadata and failure conditions for both neuron and population histories.