Indeed. This pretty much comes up whenever anyone looks at AD (Rob Trangucci being the last person to try to push on this). I had a proposal on Discourse a while back on how we could completely rebuild AD to do this. I don’t know how to do it with our current autodiff.