(author here) the goals of catgrad are a bit different to JAX - first of all, the "autodiff" algorithm is really a general approach to composing optics, of which autodiff is just a special case. Among other things, this means you can "plug your own language" into the syntax library to get a "free" autodiff algorithm. Second, catgrad itself is more like an IR right now - we're using it at hellas.ai to serve as a serverless runtime for models.

More philosophically, the motto is "write programs as morphisms directly". Rather than writing a term in some type theory which you then (maybe) give a categorical semantics, why not just work directly in a category?

Long term, the goal is to have a compiler which is a stack of categories with functors as compiler passes. The idea being that in contrast to typical compilers where you are "stuck" at a given abstraction level, this would allow you to view your code at various levels of abstractions. So for example, you could write a program, then write an x86-specific optimization for one function which you can then prove correct with respect to the more abstract program specification.