NEURAX Architecture Design

This document describes the architecture, design principles, and data flow of the NEURAX compiler system.

Table of Contents

1. Overview
2. System Architecture
3. Data Flow
4. Component Deep Dive
5. Design Principles
6. Adding a New Model Family

Overview

NEURAX is an analytical compiler for neural network architectures. Unlike traditional compilers that emit machine code, NEURAX emits a complete engineering report of a model’s behaviour on target hardware — cost, memory, speed, safety, and feasibility — in milliseconds, before a single GPU spins up.

The system is composed of:

System Architecture

graph TB
    subgraph "Frontend"
        UI[neurax-ui<br/>React 18 + TypeScript]
        TUI[neurax-tui<br/>Ratatui]
        CLI[neurax-cli<br/>Rust CLI]
    end

    subgraph "Service Layer"
        HTTP_API[neurax-service<br/>actix-web HTTP API<br/>38 REST routes]
        AI_AGENT[neurax-agent<br/>FastAPI + LangChain<br/>Natural language → architecture]
        MCP[neurax-mcp<br/>MCP server]
    end

    subgraph "Core Engine"
        CORE[neurax-core<br/>Pipeline orchestrator<br/>+ ONNX export + streaming]
        PARSER[neurax-parser<br/>JSON → ModelConfig]
        IR[neurax-ir<br/>10 IR dialects<br/>+ inference + dynamic passes]
        FORMULAS[neurax-formulas<br/>FLOPs / params / memory]
        HWDB[neurax-hardware-db<br/>20 GPUs • CPUs • interconnects]
        MLIR[neurax-mlir<br/>13 MLIR dialects<br/>LLVM 18 • IREE]
    end

    subgraph "External Services"
        SUPABASE[Supabase<br/>Auth • Database • Storage]
        STRIPE[Stripe<br/>Billing • Subscriptions]
        OPENAI[OpenAI / Anthropic<br/>LLM API]
        GITHUB[GitHub<br/>Repository • Pull Requests]
    end

    UI -- HTTP --> HTTP_API
    TUI -- direct --> CORE
    CLI -- direct --> CORE
    MCP -- HTTP --> HTTP_API

    AI_AGENT -- HTTP --> HTTP_API
    AI_AGENT -- LLM --> OPENAI

    HTTP_API --> CORE
    CORE --> PARSER
    CORE --> IR
    CORE --> FORMULAS
    CORE --> HWDB
    CORE --> MLIR

    HTTP_API -- JWT --> SUPABASE
    HTTP_API -- Billing --> STRIPE
    HTTP_API -- GitHub push --> GITHUB

    style CORE fill:#2ecc71,color:#fff
    style IR fill:#3498db,color:#fff
    style MLIR fill:#e74c3c,color:#fff
    style UI fill:#9b59b6,color:#fff
    style AI_AGENT fill:#f39c12,color:#fff

Data Flow

The primary data flow for architecture analysis:

flowchart LR
    INPUT["User Request<br/>(JSON model config)"] --> PARSER["Parser<br/>neurax-parser"]
    PARSER --> AST["ModelConfig<br/>Typed AST"]
    AST --> PIPELINE["Analytical IR Pipeline<br/>10 passes"]

    subgraph PIPELINE_CONTENT [" "]
        direction LR
        A["Arch.<br/>IR"] --> G["Graph<br/>IR"] --> T["Tensor<br/>IR"] --> O["Op<br/>IR"] --> C["Compute<br/>IR"]
        C --> M["Memory<br/>IR"] --> P["Parall.<br/>IR"] --> H["Hardware<br/>IR"] --> CO["Cost<br/>IR"] --> R["Report<br/>IR"]
    end

    PIPELINE --> REPORT["Report<br/>40+ metrics<br/>JSON / Markdown"]
    PIPELINE --> MLIR_OUT["NEURAX-MLIR<br/>model.mlir"]
    MLIR_OUT --> LLVM["LLVM 18 / IREE<br/>CPU • CUDA • ROCm<br/>Metal • Vulkan"]

    style INPUT fill:#4a90d9,color:#fff
    style REPORT fill:#2ecc71,color:#fff
    style MLIR_OUT fill:#e74c3c,color:#fff

The 10-Pass IR Pipeline

Each pass transforms the representation and computes metrics:

Dynamic Analysis (Parallel, Post-Pipeline)

Three dynamic passes run in parallel after the static pipeline:

Component Deep Dive

neurax-parser

The parser ingests JSON model configurations conforming to the NEURAX universal schema (v1.0) and produces a strongly-typed ModelConfig.

Key types:

• ModelConfig — top-level configuration (model, training, hardware, parallelism)
• ModelType enum — Transformer, CNN, MoE, SSM, Diffusion, GNN, GAN, RL, SNN, RNN, Multimodal, Custom
• LayerType enum — Attention, Mlp, Embedding, Conv2d, etc.
• Schema validation via ModelValidator

Supported model types: transformer, cnn, moe, ssm, diffusion, gnn, gan, rl, snn, rnn, multimodal, custom

neurax-ir

The IR crate implements 10 dialect-like modules, each with its own Pass struct implementing the IrPass trait:

#![allow(unused)]
fn main() {
pub trait IrPass {
    type Input;
    type Output;
    type Metrics;

    fn build(&self, input: &Self::Input, ctx: &NeuraxContext) -> Result<Self::Output, NeuraxError>;
    fn compute_metrics(&self, output: &mut Self::Output, ctx: &NeuraxContext) -> Result<Self::Metrics, NeuraxError>;
    fn validate(&self, output: &Self::Output, metrics: &Self::Metrics) -> Result<(), NeuraxError>;
}
}

Diagnostic system: Standardized diagnostic codes (E001-E005 errors, W001-W006 warnings, I001-I003 info, H001-H005 hints) with severity levels and precision impact scoring.

neurax-core

The orchestrator that wires together the 10-pass pipeline, dynamic analysis, and export. Also provides:

• run_analysis() — full pipeline entry point
• analyze_json() — JSON string → AnalysisResult
• validate_json() — JSON validation
• get_model_summary() — quick model summary
• ONNX export via neurax-core/src/export/

neurax-mlir

MLIR compiler backend with 13 custom dialects:

Lowering pipeline: Architecture → Graph → Tensor → Operator → Compute → Memory → Hardware → Cost → Report → LLVM IR → Assembly → Object file.

Target backends: CPU, CUDA, Vulkan, Metal, ROCm — plus IREE integration for cross-platform deployment.

neurax-formulas

Pure analytical formulas for ML operations. Hot path — maximum optimization.

Modules: attention, conv, mlp, embedding, normalization, moe, ssm, rnn, diffusion, gnn, custom, cnn_blocks.

neurax-hardware-db

Built-in database with 20 GPUs, 2 CPUs, and 5 interconnect specifications.

GPU specs include: H200, GH200, H100-SXM, H100-PCIe, A100-SXM, A100-PCIe, L40S, L40, V100, RTX 4090, RTX 4080, RTX 3090, RTX 6000 Ada, RTX A5000, A10G, A30, T4, K80.

Key metrics per GPU: TFLOPS (FP64/FP32/FP16/BF16/INT8/FP8), memory bandwidth, NVLink, TDP, L2 cache, SM count.

neurax-service

Production actix-web HTTP server with:

• 38 REST routes (analysis, inference, export, projects, billing, credits, compliance, API keys, agent control, presets, hardware, plugin)
• Supabase JWT authentication + API key authentication with scope-based authorization
• Stripe billing integration
• SSE streaming for real-time analysis
• CORS, gzip compression
• Health checks

neurax-agent

Python/FastAPI/LangChain AI copilot with a 3-phase declarative pipeline:

1. Planning — LLM generates a complete ArchSpec (nodes + edges) using structured output
2. Validation — Pure Python topology validator checks DAG, fan-in, connectivity
3. Materialization — Stream tool calls to the canvas with auto-correction (up to 3 retries)

Supports 11 model families with catalogues containing 400+ blocks.

neurax-ui

React 18 + TypeScript + Vite single-page application with:

• Visual canvas (React Flow) with drag-and-drop, parameter editing, minimap
• 88 reference templates across 11 families
• Metrics dashboard with 40+ metrics and charts
• AI Chat Drawer with SSE streaming
• Hyperparameter Optimization panel
• Time Machine cost/carbon projection
• Inference Intelligence panel
• Project management (cloud CRUD)
• Credits system with plan-based limits
• Export panel (ONNX, JSON, Network Graph)
• GitHub export with PR creation

neurax-mcp

Model Context Protocol server that exposes NEURAX capabilities to MCP-compatible clients (e.g., Claude Desktop). Provides 9 tools: analyze_architecture, list_templates, get_template, list_hardware, estimate_training_cost, get_compliance_config, get_credits, get_user_info, health_check.

Design Principles

1. Analytical, not empirical — All metrics are computed via pure analytical formulas. No GPU is needed, no simulation is run. Results are deterministic and available in milliseconds.

2. Compiler-inspired pipeline — The system follows the traditional compiler architecture: parse → IR → optimize → generate. Each pass is independent and composable.

3. Multi-language ecosystem — Rust for performance-critical analysis, Python for the AI agent, TypeScript for the web UI. Each language is chosen for its strengths.

4. Schema-first design — The JSON model config schema (v1.0) is the universal interchange format. All components read from and write to this schema.

5. Deterministic by default — The core analysis pipeline is fully deterministic. The AI agent uses LLMs but validates every output before materialization.

6. Extensible catalogues — Model families, blocks, and constraints are defined in JSON catalogues that can be extended without code changes.

7. Observability-first — Every analysis includes phase timing, diagnostics, and recommendations. The system explains not just what the metrics are, but why.

Adding a New Model Family

To add a new model family to NEURAX:

1. Add to the Rust parser

In neurax-parser/src/model_config.rs, add the new family to ModelType::from_str():

#![allow(unused)]
fn main() {
"my_family" => Ok(Self::MyFamily),
}

Add the corresponding serialization:

#![allow(unused)]
fn main() {
Self::MyFamily => "my_family",
}

2. Add formulas

In neurax-formulas/src/, create a new module (e.g., my_family.rs) with FLOPs, parameter, and memory formulas. Register it in lib.rs.

3. Add catalogue entries

In neurax-agent/catalogue.json, add blocks for the new family. Each block should include type, name, family, params, description, and max_inputs.

4. Add template

In templates.ts, add reference templates for the new family.

5. Add constraints

In neurax-agent/block_constraints.json, add fan-in limits for the new family’s blocks.

6. Add to arch_planner

In neurax-agent/arch_planner.py, add a family template in FAMILY_TEMPLATES describing the typical flow for the new family.