Theano Op using JAX for lightning-fast ODE inference

twiecki April 15, 2021, 3:31pm 3

This is really cool - thanks for sharing. Would be a great contribution to PyMC3!

1 Like

Topic		Replies	Views
Creating Theano Op to work on vectors element-wise within Model()? Questions	5	2352	August 30, 2017
NotImplementedError: No JAX conversion for the given `Op`: AllocDiag{offset=0, axis1=0, axis2=1} Questions	3	722	July 17, 2021
UnaryScalarOps for JAX?	3	210	May 19, 2023
Custom functions in Jax? Questions	2	815	April 5, 2021
Integration, ODE and Lambert W function in PyMC4 v5 ode_pde	11	955	December 18, 2022