A JAXPR is a normal functional-style expression tree, with free variables, like

    let b = f a in g (a, b)

Elliot's "compiling to categories" requires you to translate this to

    g ∘ (id × f) ∘ dup

It's pretty baffling to work with terms like the latter in practice! The main ideas that JAX is based on were already published in "Lambda the ultimate backpropagator" in 2008. Elliot's work is a nice way of conceptualising the AD transformations and understanding how they all relate to each other, but it's not particularly practical.