I am new to PyMC and Jax. I am trying to use PyMC to estimate the parameters for ODE implemented in Jax. The ODEs I work with often have many input parameters so I would like to write a PyTensor Op that can handle this. I have successfully used sunode, but there are some features from Jax and diffra…

I think your first problem is in vjp_sol_op_jax, have you tried to call that directly before wrapping anything in PyTensor? Are you sure the inputs should be in a tuple? Then, the Apply returned by VJPSolOp.make_node should have 4 inputs: y0, args and the two gradients, but you say it only has two …

How to write a PyTensor Op to wrap Jax ODEs with multiple input parameters

Questions v5

nlinden March 16, 2023, 4:20pm 4

Thank you! I got it working! I appreciate the help!

1 Like

Topic		Replies	Views
Theano Op using JAX for lightning-fast ODE inference Sharing	4	2831	April 26, 2021
New package: Transforming JAX to Pytensor, for ODEs and other applications Development	3	273	March 19, 2025
ODEs, PyMC4 and custom likelihood in jax Questions modeling	9	1574	April 10, 2022
How can I output a gradient in vector format in Op.grad instance? v5 jax	7	900	January 14, 2023
ODE with matrix calculation v3	3	445	April 25, 2022

How to write a PyTensor Op to wrap Jax ODEs with multiple input parameters

Related topics