mom_u_implicit_r
mom_v_implicit_r
gad_implicit_r
autodiff_store
autodiff_restore