Indeed, it doesn’t support anything other than Conv/Linear, so would need some work. In ideal world, an autovectorization operation would be able to do this automatically, but right now it seems manual algebra is required.
@richard one common demand for per-example gradient norms is from privacy people – a datapoint with a large gradient could influence the final solution strongly, so you could “leak” information about that datapoint through the final model.