@Kazuki_Osawa thanks for sharing, that is cool. Right now I am only using 2dConv / 2dTransposedConv / 2dBatchNorm / Linear layers, so perhaps it would not be too much work to assemble the components for manual Jacobian calculations. I have been changing my architecture frequently, so the ability to have things “just work” is quite nice :).
I will keep this in mind though if I need more speed.
Curious if you benchmarked JAX at all against your manual mode? At some point just implementing stuff in numpy and using JAX is going to be simpler… or going “full manual”