baileywickham/jixp

jixp is a small programming language for describing neural networks. You write model definitions in a Lisp-like syntax, then compile them into standard Python files that use JAX, a popular library for numerical computing. The generated code has no runtime dependency on jixp itself, so the language acts purely as a build step. The project is built in three layers. First, a parser reads s-expressions, the parenthesized syntax common in Lisp, and converts them into plain Python lists. Second, an evaluator provides basic Lisp features like conditionals, function definitions, and arithmetic. Third, a compiler takes model descriptions written in this language and produces a Python module with functions for initializing parameters and running the model. The compiler uses the evaluator from the second layer to compute dimension values at compile time. The example in the README shows a GPT-style language model defined in jixp. You specify dimensions like vocabulary size and layer count, then describe the model as a sequence of blocks including embeddings, attention layers, and normalization steps. Running the compiler produces a readable Python file where each defined block becomes a pair of functions, one for initialization and one for applying the model to input data. Training code imports this generated module and uses ordinary JAX functions for gradients and optimization. jixp also includes a visualization feature. Compiling with a flag produces a standalone HTML file with an interactive explorer showing the model structure as a collapsible tree, a flow diagram, and a parameter map. The tool computes exact parameter counts for each part of the model. The README includes examples for training a small network on XOR, an MNIST classifier reaching about 98 percent accuracy, and the GPT model trained on Shakespeare text. The compiler can also target MLX, a framework for running models on Apple silicon GPUs. The generated MLX code matches the JAX version numerically, with tests confirming the outputs agree. The full README is longer than what was shown.