Backend

A Backend is the adapter that lets NeuronScope observe a framework without owning it. Every backend implements the same small interface — enumerate modules, install a hook, snapshot a tensor — and every downstream primitive is framework-agnostic as a result.

Why an adapter layer

Recorders, probes, and fingerprinters never call into a framework directly. They callBackend methods. This keeps the core small, keeps framework versions from infecting the artifact format, and lets the same script record against PyTorch today and JAX tomorrow.

Shipped backends

• torch — hook-based, supports FSDP and torch.compile boundaries.
• jax — trace-based, records at jit boundaries.
• tf — graph-based, records via tf.function tracing.
• onnx — read-only, for fingerprinting exported graphs.

Registering a custom backend

import neuronscope as ns

class MyBackend(ns.Backend):
    name = "mybackend"
    def matches(self, obj) -> bool: ...
    def modules(self, obj): ...
    def install_hook(self, module, fn): ...
    def snapshot(self, tensor): ...

ns.registry.register(MyBackend())

Pitfalls

• Two backends claiming the same object — resolution is by registration order; make matches strict.
• Snapshotting without .detach() holds the autograd graph and blows memory.
• JAX backends must record inside a jit boundary or hooks are traced away.