Numerical Optimization in Aesara Graphs

I am working on a framework for solving and estimating large-scale macroeconomic models, which needs to:

  1. Compute a fixed point of a system of nonlinear equations, given parameters
  2. Solve for a linear approximation to a policy function around this fixed point
  3. Do Bayesian estimation of model parameters from data, re-computing (1) and (2) after each MCMC sampler draw.

I am feeling optimistic that it might be possible to do (3) using a NUTS sampler by implementing the whole system in Aesara, which would make the entire model solving and approximation process end-to-end differentiable.

I am working on step (1) right now, and I am curious whether there are ways to perform numerical minimization algorithms, basically any of the usual nonlinear solvers the scipy.optimize.root and .root_scalar wrap, inside an Aesara graph. I had a look at the docs, but in that context “optimization” doesn’t appear to have the meaning I hoped it would.

Any resources on this topic would be greatly appreciated!

Optimization routines can be implemented in Aesara; the necessary control flow Ops are available (e.g. aesara.scan for looping, aesara.ifelse.ifelse and aesara.tensor.switch for conditioning). See AeMCMC and AeHMC for examples of entire MCMC routines implemented in Aesara.

If you have any other Aesara-specific questions, feel free to start a discussion in the Aesara repository.