Internal Architecture & Theoretical Mechanics¶

Document Version: 0.0.1

Scope: Core Engine, Compiler Pipeline, Ontology Maintenance, and Extension Architecture

1. Abstract¶

ml-switcheroo is a deterministic, specification-driven source-to-source compiler engineered to solve the \(O(N^2)\) interoperability challenge in deep learning infrastructure. It creates a semantic bridge between high-level frameworks ( PyTorch, JAX, TensorFlow, Keras, MLX, Flax NNX, PaxML), low-level representations (MLIR, StableHLO, NVIDIA SASS, AMD RDNA), and visual domain-specific languages (TikZ, HTML).

The system’s fundamental axiom is Bidirectional Isomorphism: every supported language is treated potentially as both a Source and a Target. Whether the input is a Python AST defining a neural network or a stream of RDNA assembly instructions, the engine lifts the code into a centralized Abstract Standard (The Hub) before projecting it into the desired destination dialect (The Spoke).

This document details the internal subsystems, justifying the architectural bifurcation between the High-Fidelity Rewriter (for structured languages) and the Graph Compiler (for linear/unstructured languages), and elucidating the automated mechanisms of knowledge acquisition that keep the internal ontology synchronized with the evolving ML ecosystem using LLM-assisted loops and upstream standards ingestion.

2. The Grand Unified Architecture¶

The ASTEngine serves as the central dispatch controller. It classifies the input and output languages based on their structural properties and routes them through one of two isomorphic pipelines.

        %%{init: {'theme': 'base', 'themeVariables': { 'fontFamily': 'Google Sans', 'fontSize': '14px', 'lineColor': '#20344b'}}}%%
graph TD
    %% --- DESIGN CONSTRAINTS ---
    classDef header font-family:'Google Sans Medium',fill:#20344b,color:#ffffff,stroke:#20344b;
    classDef process fill:#4285f4,stroke:#20344b,stroke-width:1px,color:#ffffff,font-family:'Google Sans',rx:4px;
    classDef decision fill:#f9ab00,stroke:#20344b,stroke-width:1px,color:#20344b,font-family:'Google Sans Medium',rx:20px;
    classDef artifact fill:#ffffff,stroke:#20344b,stroke-width:1px,color:#20344b,font-family:'Roboto Mono',stroke-dasharray: 2 2;
    classDef storage fill:#ea4335,stroke:#20344b,stroke-width:1px,color:#ffffff,font-family:'Google Sans',rx:4px;
    classDef asm fill:#20344b,stroke:#20344b,stroke-width:1px,color:#ffffff,font-family:'Roboto Mono',rx:2px;

    SRC("Input Artifact<br/>(Source)"):::artifact
    
    subgraph ENGINE [" ⚡ ASTEngine Orchestration "]
        direction TB
        DISPATCH{"Language Topology?"}:::decision
        
        subgraph PATH_A [" 🟢 Path A: Rewriter (CST Preservation) "]
            direction TB
            INGEST_CST("<b>Ingest / Parse</b><br/>(Source &rarr; CST)"):::process
            TRANSFORM("<b>Pivot Transformation</b><br/>(CST &harr; Hub &harr; CST)"):::process
            EMIT_CST("<b>Emission</b><br/>(CST &rarr; Target)"):::process
            
            INGEST_CST --> TRANSFORM --> EMIT_CST
        end

        subgraph PATH_B [" 🔵 Path B: Compiler (Graph Synthesis) "]
            direction TB
            LIFT("<b>Lifter</b><br/>(Source &rarr; IR)"):::process
            IR("<b>LogicalGraph</b><br/>(Topological DAG)"):::artifact
            SYNTH("<b>Synthesizer</b><br/>(IR &rarr; Target)"):::process
            
            LIFT --> IR --> SYNTH
        end
    end

    KB[("<b>Semantics Manager</b><br/>(The Abstract Standard)")]:::storage

    TGT("Output Artifact<br/>(Target)"):::artifact

    %% --- EDGES ---
    SRC --> DISPATCH
    
    DISPATCH -->|" Structured (Python, MLIR, StableHLO) "| PATH_A
    DISPATCH -->|" Unstructured (ASM, TikZ, HTML) "| PATH_B

    TRANSFORM <-->|" Query Specs / Inject Plugins "| KB
    SYNTH <-->|" Query Macros "| KB

    EMIT_CST --> TGT
    SYNTH --> TGT

2.1. Path A: The High-Fidelity Rewriter¶

Ecosystem: Pure Python (PyTorch, JAX, Flax, TF, Keras, MLX, PaxML) and Textual IRs (MLIR, StableHLO).

This path treats code as a mutable abstract document. The primary objective is preservation. When transcoding high-level code, comments, whitespace, and variable names must be retained.

Mechanism: The PivotRewriter walks the tree, identifies semantic nodes (Calls, Classes), maps them to the Abstract Hub, and mutates them in-place to match the Target Spoke configuration.
Intermediate Representations: MLIR and StableHLO are treated as textual formats in this path to allow high-fidelity refactoring (structure preservation) without lossy graph compilation.

2.2. Path B: The Graph Compiler¶

Ecosystem: Hardware ISAs (SASS, RDNA) and Visual DSLs (HTML Grid, TikZ).

This path treats code as a reconstructible logic flow. Assembly languages are linear streams of instructions; they lack the hierarchical structure (classes, functions) required by Path A.

Mechanism: A Lifter parses source text (e.g., SASS Mnemonics) into a LogicalGraph Intermediate Representation (IR). The Synthesizer then reconstructs valid target code (e.g., TikZ diagrams or Python Re-implementation) from the topology.

3. The Knowledge Base (The Hub)¶

The cornerstone of the system is the Internal Abstract Standard, an ontology defining the Platonic ideal of mathematical operations. It resides in src/ml_switcheroo/semantics/.

The Hub defines what an operation is (Identity, Signature, Constraints), while the Spokes (Framework Adapters) define how it is implemented.

3.1. Construction via The LLM Feedback Loop¶

The maintenance of the Knowledge Base is semi-automated, utilizing Large Language Models (LLMs) to bridge the gap between library APIs and the formal Schema.

Introspection (ml_switcheroo suggest): The system performs runtime reflection on installed libraries.
- Command: ml_switcheroo suggest 'jax.numpy.*' or ml_switcheroo suggest 'mlx.nn.layers.*'
- Action: It scans the namespace, extracts signatures and docstrings, and generates a structured prompt.
LLM Processing: This prompt is fed to an LLM, which is tasked with mapping the discovered APIs to the Operation Definition Language (ODL). It generates a YAML file reconciling the signatures of PyTorch, JAX, TensorFlow, MLX, and Flax NNX into a single abstract definition.
Hydration (ml_switcheroo define): The resulting YAML files are injected into the system.
- Command: ml_switcheroo define my_new_ops.yaml
- Action: This hydrates the JSON storage (semantics/*.json for specifications, snapshots/*.json for mappings), effectively “teaching” the compiler new math.

3.2. Construction via Standards Ingestion¶

To ground the Abstract Standard in reality, ml-switcheroo ingests upstream specifications directly.

ONNX (Neural): The OnnxSpecImporter parses the official Operators.md from the ONNX repository, populating k_neural_net.json with layer definitions common to all frameworks.
Array API (Math): The ArrayApiSpecImporter parses the Python Consortium’s standard stubs, populating k_array_api.json with mathematically rigorous tensor operations.

4. Path A: The Rewriter Internals¶

This component executes the bidirectional translation for structured languages using LibCST.

        %%{init: {'theme': 'base', 'themeVariables': { 'fontFamily': 'Google Sans', 'fontSize': '14px'}}}%%
classDiagram
    direction TD

    class RewriterContext {
        +SemanticsManager semantics
        +SymbolTable symbols
        +List~Scope~ scope_stack
    }

    class RewriterStage {
        +visit()
        +leave()
    }

    class StructureStage {
        +leave_ClassDef()
        +leave_FunctionDef()
    }
    class Note_Struct["<b>Topology Rewriter</b><br/>Inheritance & Signatures"]

    class ApiStage {
        +leave_Call()
        +leave_Attribute()
    }
    class Note_Api["<b>Logic Rewriter</b><br/>Function dispatch"]

    RewriterStage <|-- StructureStage
    RewriterStage <|-- ApiStage
    RewriterStage *-- RewriterContext
%% Styles
    style RewriterContext fill: #4285f4, color: #fff, stroke: #20344b
    style RewriterStage fill: #fff, color: #20344b, stroke: #20344b
    style StructureStage fill: #5cdb6d, color: #20344b, stroke: #20344b
    style ApiStage fill: #57caff, color: #20344b, stroke: #20344b
    style Note_Api fill: #20344b, color: white, stroke: none, font-family: 'Roboto Mono', font-size: 10px
    style Note_Struct fill: #20344b, color: white, stroke: none, font-family: 'Roboto Mono', font-size: 10px

4.1. Adding New Python Frameworks¶

Any framework is both a source and a target. Adding support requires implementing the FrameworkAdapter protocol in src/ml_switcheroo/frameworks/.

Registry: Add a file (e.g., tinygrad.py) decorated with @register_framework.
Structural Traits: Define base classes (e.g., tinygrad.Tensor) and lifecycle methods.
Discovery: Define search_modules to opt-in to the suggest loop.
Auto-Wiring: The system automatically generates the inverse mapping. If TinyGrad maps to Hub Operation Add, the engine inherently knows how to translate Add back to TinyGrad.

4.2. Handling MLIR & StableHLO¶

These are treated as structured text. The MlirParser ingests MLIR into a custom CST. The PivotRewriter treats stablehlo.abs identical to jax.numpy.abs, allowing seamless refactoring of IR code using the same semantic database as Python code.

5. Path B: The Compiler Internals¶

This path handles linear instruction streams, enabling “Decompilation” (ASM \(\to\) Python) and Visualization ( Python \(\to\) Diagram).

        %%{init: {'theme': 'base', 'themeVariables': { 'fontFamily': 'Google Sans', 'fontSize': '14px'}}}%%
graph TD
    classDef ir fill:#f9ab00,color:#20344b,stroke:#20344b;
    classDef step fill:#4285f4,color:#ffffff,stroke:#20344b,rx:5px;
    classDef artifact fill:#ffffff,stroke:#20344b,stroke-width:1px,color:#20344b,font-family:'Roboto Mono';

    SRC_ASM("Source (Assembly/TikZ)"):::artifact
    
    subgraph LIFTER_BLOCK [" 1. Lifting (Frontends) "]
        REGEX("<b>Lexer</b><br/>(Regex Patterns)"):::step
        LIFT("<b>Lifter</b><br/>(CFG/DFG Analysis)"):::step
        
        REGEX --> LIFT
    end
    
    GRAPH("<b>LogicalGraph (DAG)</b>"):::ir
    
    subgraph SYNTH_BLOCK [" 2. Synthesis (Backends) "]
        MACRO("<b>Macro Expander</b><br/>(e.g. loops &rarr; instructions)"):::step
        EMIT("<b>Emitter</b><br/>(Text Generation)"):::step
        MACRO --> EMIT
    end

    TGT_ASM("Target (Assembly/TikZ)"):::artifact

    SRC_ASM --> LIFTER_BLOCK
    LIFT --> GRAPH
    GRAPH --> SYNTH_BLOCK
    EMIT --> TGT_ASM

5.1. Adding New ISAs (ASM)¶

To add a new ISA (e.g., a custom accelerator):

Frontend: Implement a Lifter that parses the textual assembly into LogicalNodes. It must reverse-engineer loops to infer high-level semantics (e.g., detecting a 3x3 convolution loop).
Backend: Implement a Synthesizer with a RegisterAllocator to map symbolic variables to physical registers.

5.2. Visual DSLs as Targets & Sources¶

Target (Rendering): The HtmlBackend calculates topological rank from the LogicalGraph to output CSS Grid coordinates.
Source (OCR/Diagrams): The TikzParser reads LaTeX diagram code, lifting nodes and edges into the LogicalGraph. This allows converting a paper’s architecture diagram directly into executable PyTorch code.

6. Verification Loop¶

The system ensures fidelity through the HarnessGenerator.

Constraint Satisfaction: It reads the ODL constraints (Rank, Dtype) from the JSON Hub.
Fuzzing: It uses Hypothesis to generate random tensors valid for both Source and Target frameworks.
Cross-Check: It executes both the Source implementation and the Transpiled Target, asserting np.allclose(result_src, result_tgt).

This closes the loop: Introspection (Suggest) \(\to\) Definition (Hub) \(\to\) Implementation (Spoke) \(\to\) Verification ( Fuzzer).