9 ms·

Dex is also missing user authored composable program transformations, which is one of JAX’s hidden superpowers. So not quite “JAX without limitations” — but ce

by mccoyb 2y ago

Dex is also missing user authored composable program transformations, which is one of JAX’s hidden superpowers. So not quite “JAX without limitations” — but certainly without some of the limitations.

funks_ 2y ago
Are you talking about custom VJPs/JVPs?
- mccoyb 2y ago
  No, I'm talking about custom `Jaxpr` interpreters which can modify programs to do things.
6gvONxR4sf7o 2y ago
This is both its strength and its weakness. As soon as you write a jaxpr interpreter, you lose all the tooling that makes the python interpreter so mature. For example stack traces and debugging become black holes. If jax made it easy to write these transformations without losing python’s benefits it would be incredible.