Imagine if every open-source language model on the planet — Llama, Mistral, Qwen, all of them — depended on a single company's software library to even load. If that library disappeared tomorrow, a huge chunk of the AI ecosystem would grind to a halt. That's roughly the situation today: the vast majority of open models are defined in code that targets one specific Python framework (Hugging Face's transformers), and porting a model to run efficiently on different hardware or in a different training system is a painful, error-prone manual job.

This paper introduces Axon, a small programming language purpose-built for describing neural network architectures — particularly large language models — in a way that isn't chained to any single framework.

The core ideas:

• Write the model once, run it anywhere. An Axon description of, say, a transformer block can be compiled to PyTorch for training, to a fast inference engine for deployment, or to something else entirely. The model definition is decoupled from the runtime.
• Shape-safe by construction. Anyone who has debugged deep learning code knows the pain of tensor shape mismatches — matrix A is (batch, 512, 768), matrix B is (batch, 768, 512), and you get a runtime error twenty minutes into training. Axon uses a strong type system (Haskell-inspired) that checks tensor shapes before the code ever runs. Whole categories of stupid bugs simply become impossible.
• A domain-specific language, not a general one. Rather than being another Python library, Axon is its own small language designed specifically for the vocabulary of neural nets: layers, attention heads, embeddings, residual connections. That focus lets the compiler be much smarter about optimization and verification.

The why matters more than the what. Right now, when a researcher invents a new attention mechanism or a hardware team builds a new accelerator, everyone has to reimplement the same models over and over, and each reimplementation risks subtle bugs that silently degrade quality. Axon tries to break this cycle by making the model architecture a portable, verifiable artifact — closer to a mathematical specification than a piece of PyTorch code.

The key insight is treating model architectures as a compilation problem rather than a scripting problem. Once you have a typed intermediate representation, you unlock the same benefits compilers give traditional software: portability across targets, catch-bugs-at-compile-time safety, and machine-checkable optimizations.