At the moment, the ReverseDiff pullback is a bit hacky because it takes the gradient of a closure. Ideally we should hook into ReverseDiff internals