Should pullback a single contangent through the entire desc compute infrastructure to extend factor N memory and performance improvement, where N is number of rows in Jacobian, to full pipeline.
All you have to do is put the sparse pullback decorator over that single function. Few lines changed . For meshgrids this is easy.
This would require the inputs and outputs the to the registered compute functions to be of shape (rho, theta, zeta, ...) . That's actually the natural shape that np.meshgrid(rho,theta,zeta, ij) returns. Then wrap the full compute pipeline in a single sparse pullback decorator except for Zernike transform. Edit: see below for simpler
Before exiting to return data to users use tree_map(meshgrid_reshape, data) so that retuned data to users is same structure as before.
Should pullback a single contangent through the entire desc compute infrastructure to extend factor N memory and performance improvement, where N is number of rows in Jacobian, to full pipeline.
All you have to do is put the sparse pullback decorator over that single function. Few lines changed . For meshgrids this is easy.
This would require the inputs and outputs the to the registered compute functions to be of shape (rho, theta, zeta, ...) . That's actually the natural shape that
np.meshgrid(rho,theta,zeta, ij)returns. Then wrap the full compute pipeline in a single sparse pullback decorator except for Zernike transform. Edit: see below for simplerBefore exiting to return data to users use tree_map(meshgrid_reshape, data) so that retuned data to users is same structure as before.