at.grad throws NotImplementedError

[jessegrabowski]

On this episode of "If You Give a Mouse a Cookie", I have put everything you showed me in the other thread together into one system, and I am able to get gradients with respect to model parameters through two iterative solution algorithms, which is really neat! But when I try to tack on a Kalman Filter to the end, suddenly the wheels fall off the cart. Asking for gradients with respect to the parameters that go through the entire solver process throws a "NotImplementedError". If I ask for gradients with respect to variables that do NOT go through this solver, on the other hand, I get back happy gradients.

Evidently there is an Op buried somewhere in the pile that doesn't have a grad function defined, but I have no idea how to troubleshoot which one it might be. Any good tips for figuring out where the chain of gradients is breaking? I tried looking at aesara.dprint, but I didn't get anything from it.

Worth nothing is that ran into this problem once when using at.linalg.inv, but switching to at.linalg.solve fixed it previously. Also, this is the same Kalman filter code I use in my statespace package, which DOES give me gradients, so I am worried that I am just flying too close to the sun -- too many scans on top of scans on top of scans?

I uploaded a short notebook that shows the error, along with some checks that things aren't completely broken in other places.

[ricardoV94]

We should print a more informative exception that mentions which Op has no grad implemented. For now setting a breakpoint before the line that raises should allow you to find which Op it is

[jessegrabowski]

Cracked it, still had a linalg.inv floating in the code. Stupid. Sorry for spamming the board.

[ricardoV94]

If you want to take a chance at improving the Error message in the codebase that would be great!

[jessegrabowski]

Yeah I would really like to take a crack at contributing something. I'll see what I can do. You will have to hold my hand with the submission process, I've never tried to contribute to a big boy codebase before.

[ricardoV94]

No problem we are here to help. The first step would be to open an issue to describe the problem, usually with a minimal code example. This issue will then be closed once the solution is merged.

[brandonwillard]

There is a framework and messaging in aesara.grad that prints the Op when a gradient isn't available (e.g. grad_undefined). It seems like we should consider using something like that instead of NotImplementedError as the default Op.grad implementation.