Great work! Here's some prior art in the (torch) space: https://github.com/vccimaging/DiffOptics

A few notes, though paraxial approximations are "dumb", they are very useful tools for lens designers and understanding/constraining the design space - calculating the F/#, aperture stop, principal planes and is critical in some approaches. This pushes what autodiff tools are capable of because you need to get Hessians of your surface. There's also a rich history in objective function definition and quadrature integration techniques thereof which you can work to implement, and you may like to have users be able to specify explicit parametric constraints.