Optimize Jax usage in model#65

maxisi wants to merge 2 commits intomainmaxisi/ringdown:mainfrom

jax-optimizingmaxisi/ringdown:jax-optimizingCopy head branch name to clipboard