More broadly, I’d say Julia has more potential than Python when it comes to automatic differentiation. JAX is really cool, but Julia’s automatic differentiation tools are targeting the whole language without modification and restriction (outside those imposed by calculus). Potentially some really fancy stuff could get embedded in neural nets beyond differential equations and with relatively little effort.