Onnx Doctor

ONNX Doctor — Development Skills

Package Structure

• Package: onnx_doctor at src/onnx_doctor/
• Core modules: _rule.py, _rule_registry.py, _loader.py, _checker.py, _cli.py, _formatter.py, _message.py, _diagnostics.py
• Providers: diagnostics_providers/ — each provider is a subpackage or module
• Tests: src/onnx_doctor/tests/
• Docs: docs/ (Sphinx + MyST markdown, furo theme)

Rule Numbering Convention

• Spec rules that apply to both protobuf and IR use ONNX001–ONNX099.
• IR-only rules (e.g., duplicate Value object identity) use ONNX101+.
• Protobuf-only rules (impossible in IR by construction) use PB prefix.
• Simplification rules (unused functions/nodes/opsets) use SIM prefix.

Adding a New Rule

1. Define in YAML (spec.yaml or provider-specific YAML):

   - code: ONNX036
     name: kebab-case-name
     category: spec
     severity: error
     message: Short description of the issue.
     suggestion: How to fix it.
     explanation: |
       ## Details
       Extended markdown explanation.
   

2. Implement check in the provider (e.g., onnx_spec/__init__.py):

   # ONNX036: kebab-case-name
   if condition:
       yield _emit(_rule("ONNX036"), "node", node, message=f"...")
   

3. Add a test in tests/test_onnx_spec_provider.py:

   def test_kebab_case_name(self):
       model = _make_model(...)
       messages = _diagnose(model)
       self.assertIn("ONNX036", _codes(messages))
   

Build & Test

pip install -e .                              # Editable install
pip install -r requirements-dev.txt           # Dev dependencies
python -m pytest src/onnx_doctor/tests/       # Run tests
ruff check src/                               # Lint
ruff format src/                              # Format
onnx-doctor check model.onnx                  # CLI

Code Style

• Every .py file must start with from __future__ import annotations.
• Google-style docstrings. Target Python 3.9.
• Ruff enforced (see pyproject.toml for full config).
• Private modules prefixed with _ (e.g., _rule.py, _checker.py).

Key Dependencies

• onnx_ir: The linter operates on IR objects (ir.Model, ir.Graph, etc.), not protobuf directly.
• onnx: Used for op schema lookups (onnx.defs.get_schema).
• pyyaml: Rule definitions loaded from YAML files.
• rich: CLI output formatting.

Architecture Notes

• _checker.py provides the diagnose(model, providers) entry point that calls each provider's diagnose(model) method.
• Each provider is responsible for walking the model structure as needed (using ir.traversal.RecursiveGraphIterator or manual iteration).
• Location inference: The driver builds a location map by walking the model once. For messages without a location set, it infers the location from the target object (e.g., graph:node/3(MatMul)).
• _loader.py has a lazy singleton get_default_registry() that loads all YAML rule files on first access.
• Providers yield DiagnosticsMessage objects with target and target_type fields for context.

Autofix Architecture

• Fix = Callable[[], None] — a no-arg callable that mutates the IR in place. Stored on DiagnosticsMessage.fix.
• Rules marked fixable: true in YAML should attach a fix callable via the _emit() helper.
• CLI --fix applies all fixes, saves the model, then re-diagnoses to show remaining issues.
• CLI --diff shows a unified diff of what --fix would change, without writing.
• Fix deduplication: _apply_fixes() deduplicates by callable identity (id(fix)) to avoid running the same pass multiple times.

Available IR Passes for Fixes

From onnx_ir.passes.common (all take model: ir.Model, return PassResult):

Adding a Fixable Rule

1. Mark fixable: true in YAML.

2. In the provider, pass fix= to _emit():

   yield _emit(
       _rule("ONNX004"), "graph", graph,
       fix=graph.sort,
   )
   

3. For model-level passes, capture the model in a closure:

   yield _emit(
       _rule("ONNX003"), "graph", graph,
       fix=lambda: _apply_name_fix(model),
   )
   

Provider Structure

Providers implement a single diagnose(model: ir.Model) method that yields DiagnosticsMessage objects. Each provider is responsible for its own traversal strategy.

Example Provider Implementation

class MyProvider(onnx_doctor.DiagnosticsProvider):
    def diagnose(self, model: ir.Model) -> onnx_doctor.DiagnosticsMessageIterator:
        # Model-level checks
        if not model.graph.name:
            yield _emit(_rule("MYRULE001"), "graph", model.graph)

        # Walk all nodes (including subgraphs)
        for node in ir.traversal.RecursiveGraphIterator(model.graph):
            if some_condition(node):
                yield _emit(_rule("MYRULE002"), "node", node)

        # Check functions
        for func in model.functions.values():
            yield from self._check_function(func)

Converted and distributed by TomeVault — claim your Tome and manage your conversions.