GilesStrong / tomopt

TomOpt: Differential Muon Tomography Optimisation
GNU Affero General Public License v3.0
4 stars 0 forks source link

Improve: move to functorch vmap #172

Closed GilesStrong closed 1 year ago

GilesStrong commented 1 year ago

Swaps out pytorch internal vmap for functorch vmap: same functionality, but more performant implementation. Adds functorch as dependency.