talon-bridge 0.0.2

TALON Bridge: PyTorch export, import, and runtime for Talon IR graphs

talon-bridge - PyTorch Bridge for TALON IR

PyTorch helpers for the TALON IR (Tactical AI at Low-power On-device Nodes Intermediate Representation). This package enables PyTorch-based libraries to translate models to and from TALON IR, providing export, import, and execution capabilities for neuromorphic neural networks.

talon-bridge serves as the bridge between PyTorch models and TALON IR graphs, enabling seamless model portability and deployment to TALON hardware.

Read more about TALON IR in our documentation

Installation

uv add talon-bridge@git+ssh://[e-mail]:type1compute/talonbridge.git

Or with pip:

pip install talon-bridge

Usage

talon-bridge provides two main functions to_ir and ir_to_torch that map PyTorch models to TALON IR graphs and vice versa.

Translating from PyTorch to TALON IR

talon-bridge automatically traces through your PyTorch model and converts each module to the corresponding TALON IR primitive. You can provide custom converters for unsupported modules via the custom_map argument.

import torch
import torch.nn as nn
from talon import bridge, ir

# Define your model
model = nn.Sequential(
    nn.Linear(784, 128),
    nn.Linear(128, 10)
)

# Export to TALON IR
graph = bridge.to_ir(model, torch.randn(1, 784))

# Save to file
ir.write("model.t1c", graph)

For custom modules, provide a mapping dictionary:

# Define custom PyTorch -> TALON IR mapping
from talon import bridge, ir
custom_map = {}
def _extract_custom_module(module: MyCustomModule) -> ir.Node:
    # Convert your custom module to a TALON IR node
    return ir.Affine(weight=module.weight, bias=module.bias, ...)
custom_map[MyCustomModule] = _extract_custom_module

# Export with custom mapping
graph = bridge.to_ir(model, sample, custom_map=custom_map)

Translating from TALON IR to PyTorch

talon-bridge automatically creates PyTorch modules from TALON IR nodes. The ir_to_torch function combines import with execution, returning an executable GraphExecutor module.

from talon import bridge
import torch

# Load and create executable module
executor = bridge.ir_to_torch("model.t1c")

# Execute
output = executor(input_tensor)

For stateful models (SNNs), use return_state=True:

# Load with state management
executor = bridge.ir_to_torch("model.t1c", return_state=True)

# Execute with state
state = None
for t in range(100):
    output, state = executor(input_tensor, state)

You can also provide custom node converters:

from talon import bridge, ir
def custom_node_map(node: ir.Node) -> Optional[torch.nn.Module]:
    if isinstance(node, MyCustomNode):
        return MyCustomModule(...)
    return None  # Falls back to default converter

executor = bridge.ir_to_torch(graph, node_map=custom_node_map)

Quick Start

Basic Export and Execute

import torch
import torch.nn as nn
from talon import bridge, ir

# Your model
model = nn.Sequential(nn.Linear(784, 128), nn.Linear(128, 10))
sample = torch.randn(1, 784)

# Export to TALON IR
graph = bridge.to_ir(model, sample)

# Save to file
ir.write("model.t1c", graph)

# Load and execute
executor = bridge.ir_to_torch("model.t1c")
output = executor(sample)

Spiking Neural Network

import torch
import torch.nn as nn
import snntorch as snn
from talon import bridge

# Define SNN
model = nn.Sequential(
    nn.Linear(784, 128),
    snn.Leaky(beta=0.9, init_hidden=True, output=True),
    nn.Linear(128, 10),
    snn.Leaky(beta=0.9, init_hidden=True, output=True),
)

# Export
graph = bridge.to_ir(model, torch.randn(1, 784))

# Execute with state
executor = bridge.ir_to_torch(graph, return_state=True)
state = None
for t in range(100):
    output, state = executor(input_tensor, state)

Supported Modules

talon-bridge automatically converts the following PyTorch modules to TALON IR:

PyTorch Module	TALON IR Node	Notes
nn.Linear	Affine or SpikingAffine	Use spiking=True for hardware-optimized version
nn.Conv2d	Conv2d	Full 2D convolution support
nn.Flatten	Flatten	Tensor flattening
nn.MaxPool2d	MaxPool2d	2D max pooling
nn.Identity	Skip	Skip/residual connections
snn.Leaky	LIF	snnTorch LIF neuron

For Hardware Teams

talon-bridge is for validation only. The execution runtime is a PyTorch simulation of what the hardware will do. Use it to verify your IR graph produces correct outputs before deploying to actual hardware.

The GraphExecutor handles:

• Topological sorting of graph nodes
• State management for stateful modules (LIF neurons)
• Skip connections and residual paths
• Multiple input/output handling

Development

From the repo root, install in editable mode (requires talon-ir installed first):

uv pip install -e .
# or: pip install -e .
pytest tests/ -v

If you see AttributeError: module 'talon.bridge' has no attribute 'to_ir', you are likely using an old install. Uninstall any legacy package and reinstall in editable mode:

pip uninstall talon-bridge 2>/dev/null; uv pip install -e .

Documentation

For more information, visit the TALON IR documentation.