The authors gives that quote (from the JAX documentation) but does not seem to interiorize it as his conclusion says: > This is the niche that Theano (or rathe

The authors gives that quote (from the JAX documentation) but does not seem to interiorize it as his conclusion says: > This is the niche that Theano (or rather, Theano-PyMC/Aesara) fills that other contemporary tensor computation libraries do not: the promise is that if you take the time to specify your computation up front and all at once, Theano can optimize the living daylight out of your computation - whether by graph manipulation, efficient compilation or something else entirely - and that this is something you would only need to do once. It is exactly what JAX does. There is a computational graph in JAX (its encoded in XLA and specified with their numpy like syntax), it is build once, optimized and then runs on the GPU.