Question: status of multi-node multi-GPU parallel jax-cfd

Hi jax-cfd team!
Thanks for the great work here!

I wanted to ask folks here about the current plan regarding the extension of jax-cfd to multi-GPU systems. We're interested in using jax-cfd in a multi-host multi-GPU setting, but it seems that at the moment this is not yet supported. Our team would potentially be interested in developing such an extension, but I wanted to first coordinate with developers here what the status of this topic was.

Thanks!