Jixp, a Lisp DSL for describing Jax neural nets

Hacker News

Share card

Jixp, a Lisp DSL for describing Jax neural nets

This is a side project I've been working on while learning Jax. I noticed that a bunch of the neural net math looked like it would work well in a lisp syntax because most of the data flows through layers in a "functional" manner. Data is threaded through one layer at a time, each layer composing with the previous layer. Jixp is designed as a learning tool for myself while learning Jax to toy around with different shapes/styles of models. I used this to train a 4m parameter model on my Obsidian vault to see if I could use it for recall which partially worked. ```lisp (let-dim (d 256) (heads 8) (define transformer-block (chain (residual (layernorm d) (attention d heads #:causal)) (residual (layernorm d) (mlp d [4d] d #:bias)))) ``` The Jixp compiler reads the (racket style) DSL and outputs python which is then consumed via your normal python training stack. The advantage here is you get a pretty simple interface to describe your model and the implementation details are largely hidden in the generated python. I definitely do not recommend using this for anything besides learning. One day I can see models being specified in terms of a DSL like this which would allow different implementations of inference/training to load the same model. You could write one model definition, then vLLM, PyTorch, and your custom inference stack could all use the same definition. Sort of like a `.safetensor` file for the model.