Make jaxsplat work on modern jax.#4

OzanCatalVerses wants to merge 3 commits intoyklcs:mainyklcs/jaxsplat:mainfrom

VersesTech:mainVersesTech/jaxsplat:mainCopy head branch name to clipboard