Possible standalone distribution of `xarray_jax`?

Hi! I'm a big fan of the boilerplate code used to wrap `xarray` into a JAX-compatible entity -- I think this could have wider usage were it more well-known, especially for this kind of deep learning + weather project.

Would you consider distributing this code as a stand-alone helper module? I'm happy to volunteer refactoring this code into a small library, since I'm probably going to use it anyway -- let me know if you'd like me to take the initiative on that (and where it should live).