Aidge interop torch API#

Introduction #

aidge_interop_torch provides interoperability between Aidge and PyTorch by wrapping a torch.nn.Module into an Aidge graph (via ONNX export/import) so that the two frameworks can run side by side on the same weights and inputs. The test suite documented here does not test business features — it validates the numerical and performance equivalence between a PyTorch model and its Aidge counterpart, for both the forward pass and the backward pass (gradients and loss).

Every test in the suite follows the same pattern:

Build a PyTorch model.
Convert it into an Aidge model with aidge_interop_torch.wrap().
Run identical inputs through both models.
Compare outputs, loss, gradients and (optionally) wall-clock time.

The aidge_interop_torch.wrap() function exports the PyTorch model to ONNX, imports it into Aidge, and copies the model parameters. This ensures that both frameworks execute the same computation graph with identical weights, making numerical comparisons meaningful.

Basic example of a comparison, for a single step (see ModelTester for the full multi-epoch driver):

import torch
import aidge_interop_torch
from utils import ModelTester

torch_model = torch.nn.Linear(16, 4)
aidge_model = aidge_interop_torch.wrap(torch_model, [(1, 16)])

tester = ModelTester(
    model1=torch_model,
    model2=aidge_model,
    name="Linear_example",
    test_backward=True,
    epochs=10,
)
tester.test_multiple_step(input_sizes=(1, 16), label_sizes=(1, 4))

Three test suites are built on top of this pattern:

test_models.py — full, real-world networks (LeNet, MobileNet, ResNet, …).
test_operators.py — one operator at a time (Conv2D, ReLU, BatchNorm, …), parametrized from JSON benchmark configurations.
test_tensor.py — round-trip conversion of raw tensors between the two frameworks (dtype, shape, CPU/GPU).

A results-analysis script, analyze_results.py, post-processes pytest output to classify failures and summarize performance.

Components #

ModelTester #

ModelTester (defined in utils.py) is the comparison engine used by every test suite. It owns both models, runs matched forward/backward passes, and asserts that outputs, losses and gradients agree within a tolerance. It can optionally record forward/backward timings and fail the test if Aidge is significantly slower than PyTorch.

class aidge_interop_torch.ModelTester(model1, model2, name='', test_backward=True, epochs=10, relative_precision=0.001, absolute_precision=0.0001, learning_rate=0.01, cuda=False, benchmark=True, max_slowdown_ratio=10.0, bench_warmup=2, bench_print=True, bench_min_time_ms=0.1)#

Parameters:

model1 (torch.nn.Module) – the reference PyTorch model.
model2 (torch.nn.Module) – the Aidge model to validate (typically the output of aidge_interop_torch.wrap()).
name (str) – label used in assertion messages and benchmark records.
test_backward (bool) – if True, also runs a backward pass, compares the loss, input gradients and parameter gradients, and updates model1 with an SGD step.
epochs (int) – number of forward/backward iterations performed by test_multiple_step().
relative_precision (float) – relative tolerance used when comparing tensors.
absolute_precision (float) – absolute tolerance used when comparing tensors.
learning_rate (float) – learning rate used during backward tests.
cuda (bool) – run both models on GPU and synchronize before every timing measurement.
benchmark (bool) – enable forward/backward timing collection.
max_slowdown_ratio (float) – maximum tolerated mean(aidge) / mean(torch) ratio before test_multiple_step() raises.
bench_warmup (int) – number of leading epochs excluded from timing (JIT / cache warm-up).
bench_print (bool) – print a one-line benchmark summary at the end of test_multiple_step().
bench_min_time_ms (float) – skip the slowdown-ratio assertion for a phase whose PyTorch time is below this threshold, since the measurement is overhead-dominated at that scale.

unit_test(input_tensors, label, record_timing=False)#: Run a single forward pass (and, if test_backward, a backward pass) on both models with the same input tensors, then assert that outputs, loss and gradients match. Raises AssertionError on mismatch.

test_multiple_step(input_sizes, label_sizes=None)#: Run unit_test() for epochs iterations, generating fresh random inputs (and, for the backward pass, all-ones labels) at each step. Wraps any AssertionError with the failing epoch number. At the end, prints and/or logs the benchmark summary if benchmark=True.

performance_stats()#: Return a dict with "forward" and "backward" entries, each containing mean1 (PyTorch), mean2 (Aidge), ratio and samples. Times are in seconds.

BENCH_LOG_FILE: str | None = "benchmarks.jsonl"#: Class attribute. Path of a JSONL file that benchmark records are appended to (one JSON object per test: name, forward, backward). Written to disk rather than printed because pytest captures stdout by default. Set to None to disable file logging.

Operator test framework #

conftest.py provides a small framework for turning an operator plus a JSON benchmark configuration into a fully parametrized pytest fixture. It lets each operator declare, in a few lines, where its test cases come from and how its JSON attributes map onto PyTorch constructor kwargs.

class aidge_interop_torch.OperatorCase#

One parametrized test case, as consumed by the fixture built by make_operator_fixture().

Parameters:

id (str) – the pytest parameter id (shown in -k filters).
op_class (type) – the torch.nn.Module subclass to instantiate.
torch_kwargs (dict) – keyword arguments passed to op_class(...).
input_size (tuple) – shape of the input tensor.
output_size (tuple) – shape of the output tensor, or None to reuse input_size.
test_backward (bool) – whether to run and compare the backward pass.
expected_failure (str) – if set, the case is registered as pytest.mark.xfail with this reason instead of failing the suite.

class aidge_interop_torch.OperatorConfig#

Full description of one operator under test: either a JSON file to parse, or a hand-written list of OperatorCase.

Parameters:

name (str) – operator name, used as a prefix for generated test ids.
op_class (type) – the torch.nn.Module subclass under test.
json_path (pathlib.Path) – path to a benchmark JSON file (see JSON configuration formats below). Ignored if manual_cases is set.
manual_cases (list) – an explicit list of OperatorCase, bypassing JSON parsing entirely.
supported_attrs (dict) – mapping of JSON attribute name to (torch_kwarg_name, type), defaults to the module-level SUPPORTED_ATTRS table (kernel/stride/dilation/padding, axes, alpha, beta, …).
ignored_attrs (frozenset) – JSON attribute names to drop for this operator (e.g. an Aidge-only attribute with no PyTorch equivalent).
skip_reasons (dict) – mapping of JSON case key to an xfail reason.
default_input_size (tuple) – fallback input shape used when a JSON case does not specify one.
specs_to_kwargs (callable) – (merged_input_specs) -> dict, used to derive constructor kwargs that depend on tensor shapes (e.g. in_channels/out_channels for Conv2D or FC).
xfail_all (str) – if set, every case for this operator is xfailed with this reason.
skip_all (str) – if set, every case is skipped outright (for cases that crash the interpreter, where xfail cannot recover).
xfail_backward (str) – if set, only the backward-testing cases are xfailed; forward-only cases still run normally.
forward_only_cases (bool) – whether to keep forward-only JSON cases (shape sections); defaults to the module-level INCLUDE_FORWARD_ONLY_CASES toggle.
backend (str) – force "cpu" or "cuda" for this operator.

aidge_interop_torch.build_cases(cfg)#

Parameters:: cfg (OperatorConfig)
Returns:: the flattened, filtered list of OperatorCase for cfg (JSON-parsed or manual), with forward-only cases removed unless requested.
Return type:: list[OperatorCase]

aidge_interop_torch.make_operator_fixture(cfg)#

Build and return a pytest.fixture parametrized over every case produced by build_cases(cfg). Each parameter wraps the corresponding PyTorch model with aidge_interop_torch.wrap(), attaches an Aidge SGD optimizer when the case trains, and applies the right xfail/skip mark (operator-wide skip, operator-wide xfail, backward-only xfail, or per-case xfail, in that priority order).

Parameters:: cfg (OperatorConfig)
Returns:: a fixture yielding a dict with keys name, torch_model, aidge_model, input_size, output_size, test_backward.

JSON configuration formats. build_cases automatically detects one of the following JSON layouts:

Format	Structure
`benchmark_ok`	`base_configuration` + `test_configurations.ok.{attr,shape}`
`benchmark_flat`	`base_configuration` + `test_configurations.{shape,attributes,sweeps,…}`
`simple`	no `base_configuration`; flat, named cases
`manual`	no JSON; `OperatorConfig.manual_cases` supplies the list

For operators whose constructor needs dimensions that live inside the JSON (Conv2D.in_channels, FC.out_features, …), supply a specs_to_kwargs extractor:

def my_extractor(merged_specs: list[dict]) -> dict:
    # merged_specs = base input_specs overridden by per-case input_specs
    return {"in_channels": merged_specs[0]["shape"][1]}

MY_OP_CONFIG = OperatorConfig(
    name="MyOp",
    json_path=CONFIG_ROOT / "my_op.json",
    op_class=MyOpModel,
    specs_to_kwargs=my_extractor,
)

Test suites #

test_models.py #

Runs ModelTester on complete networks — anomaly_model, kws_model, lenet_model, mobilenet_model, perceptron_model, resnet_model — in both eval-only and train configurations. Cases are declared as a plain list of dicts (test_cases) and turned into a single parametrized fixture, model_test_case; known-failing configurations carry an expected_failure string and are marked xfail automatically.

pytest test_models.py -v -k "LeNet_eval"

test_operators.py #

Runs ModelTester operator-by-operator: pooling, activations, Conv2D/ConvDepthwise2D/ConvTranspose2D, FC, reductions, elementwise Add/Mul/Sub, BatchNorm, LayerNorm, Pad, and Aidge-specific ops (ShiftGeLU, ShiftMax). Each operator is declared once as an OperatorConfig — either JSON-backed or with manual_cases for operators with no benchmark JSON — then turned into a fixture with make_operator_fixture() and driven by a small test_<Op>_interop function.

pytest test_operators.py -v                                # all operators
pytest test_operators.py -v -k "Conv2D"                     # one operator
pytest test_operators.py -v -k "Conv2D_attributes_kernel_square"
pytest test_operators.py -v -k "ReduceMean_axes_1_dim"

Adding a new operator:

Create operators/<op>.py (a single nn.Module).
Export it from operators/__init__.py.
Declare an OperatorConfig (with json_path or manual_cases).
Call make_operator_fixture() and write a test_<Op>_interop test.

test_tensor.py #

Round-trips raw tensors — not models — through aidge_interop_torch.tensor_converter.torch_tensor_to_aidge() and aidge_interop_torch.tensor_converter.aidge_tensor_to_torch(), and checks values, dtypes and shapes are preserved. Covers standard dtypes (float32/float64/int32/int64), several shapes including empty and scalar tensors, and CPU/GPU transfer (GPU cases are skipped when CUDA or aidge_backend_cuda is unavailable).

pytest test_tensor.py -v -k "test_conversion_cpu"

Results analysis #

analyze_results.py parses saved pytest output (verbose or quiet) and the benchmarks.jsonl log produced by ModelTester.BENCH_LOG_FILE, then prints a categorized report: Aidge bugs to report, performance regressions, test-strategy issues (e.g. shape mismatches), and a per-operator benchmark summary with the ten worst forward/backward slowdown ratios.

pytest test_operators.py -v 2>&1 | tee results.txt
python analyze_results.py results.txt

# explicit benchmark file
python analyze_results.py results.txt --bench benchmarks.jsonl

# or piped directly
pytest test_operators.py -v 2>&1 | python analyze_results.py

aidge_interop_torch.categorize(rt_msg, ae_msg)#

Parameters:

rt_msg (str) – the RuntimeError message captured from a failure, if any.
ae_msg (str) – the AssertionError message captured from a failure, if any.

Returns:

a (category, emoji) pair, e.g. ("outputs_differ", "🔴") or ("shape_mismatch (test strategy)", "🟡").

Return type:

tuple[str, str]

Adding a new interoperability test #

The interoperability framework is designed so that adding support for a new operator or model requires very little boilerplate. Most of the comparison logic is handled by ModelTester and the fixtures generated by make_operator_fixture().

Model tests #

Model interoperability tests validate complete neural networks instead of individual operators.

To add a new model:

Create the PyTorch model

Add a new model implementation under models/. The model should inherit from torch.nn.Module and implement the desired network architecture.

For example, a LeNet implementation can be added as:

class TorchLeNet(torch.nn.Module):

    def __init__(self):
        ...
        self.sequence = torch.nn.Sequential(...)

    def forward(self, x):
        return self.sequence(x)

Register the model

Add a new entry to the test_cases list in test_models.py.
```
{
    "name": "LeNet_train",
    "model": lenet_model.TorchLeNet(),
    "input_size": (BATCH_SIZE, 1, 32, 32),
    "output_size": (BATCH_SIZE, 10),
    "test_backward": True,
    "expected_failure": "Backward of BatchNorm is not implemented",
}
```
Each model configuration defines:
- a unique test name;
- the instantiated PyTorch model;
- the input and output tensor shapes;
- whether backward propagation should be validated;
- an optional expected failure for known limitations.

Once registered, the model is automatically included in the parametrized pytest suite and compared against its Aidge counterpart using ModelTester.

Operator tests #

Adding a new operator generally consists of three steps.

Create a PyTorch wrapper

Create a new file under operators/ containing a minimal torch.nn.Module that exposes the operator to test.

For example, an AvgPool2d wrapper can be written as:

class AvgPoolingModel(nn.Module):

    def __init__(
        self,
        kernel_size=(3, 3),
        stride=(2, 2),
        padding=(0, 0),
        ceil_mode=False,
        count_include_pad=True,
    ):
        super().__init__()
        self.pool = nn.AvgPool2d(
            kernel_size=kernel_size,
            stride=stride,
            padding=padding,
            ceil_mode=ceil_mode,
            count_include_pad=count_include_pad,
        )

    def forward(self, x):
        return self.pool(x)

Declare an OperatorConfig

Register the operator in test_operators.py by creating an OperatorConfig describing how the operator should be tested.

AVGPOOLING_CONFIG = OperatorConfig(
    name="AvgPooling2D",
    json_path=CONFIG_ROOT / "avgpooling2d.json",
    op_class=AvgPoolingModel,
    ignored_attrs=frozenset(
        {
            "rounding_mode",  # Aidge-specific
            "dilations",      # not used by AvgPool
        }
    ),
    backend="cuda",
)

The configuration specifies:

the operator name used in generated test identifiers;
the JSON benchmark configuration describing the test cases;
the PyTorch wrapper class;
optional ignored or remapped attributes;
backend-specific restrictions, when applicable.

Provide the benchmark configuration

Finally, add the corresponding JSON benchmark file under the benchmark configuration directory. The interoperability framework automatically generates the pytest cases from this configuration using build_cases() and make_operator_fixture().