Add docs on interfacing with surrogates

[theo-brown]

Includes:

1. Manually reimplementing the model in JAX
2. Converting a Pytorch model to a JAX model using torch_xla2
3. Using an ONNX model with jaxonnxruntime

Potentially closes #538

[jcitrin]

Thanks for this! Really awesome and appreciated that you are writing documentation.

Made a few comments. @hamelphi , @sbodenstein , @ernoc could you also take a look at this?

[jcitrin]

When fusion_transport_surrogates matures more it will hopefully help abstract away some of this and we can supplement/modify these docs with examples using that library

[hamelphi]

Nice! LGTM. Thanks for the contribution Theo.

[theo-brown]

Thanks for the comments, sorry for the delay in responding!

[theo-brown]

Following some of the suggestions made by @sbodenstein, I've added a bit on saving/loading models in HLO format, which is the one backed by OpenXLA.

[theo-brown]

I've just made the changes @sbodenstein requested. However, after rebasing onto main I've discovered that the relevant docs page has changed to no longer contain information on how to integrate new models into TORAX - I imagine because a specific API/framework for doing so is to be developed, and you guys are currently crunching for v1.0.0. Is this PR dead, then, or will this info be useful somewhere sometime in the future?

[jcitrin]

Hi Theo. I think it's still useful to include and if done right now can be part of the v1.0.0 docs (release planned imminently). Best would be to not include all of interfacing_with_surrogates.rst into the near-empty model_integration.rst, but to add instead a link with some text like "For information on JAX-friendly interfacing with ML-surrogates of physics models, see <your_new_link>".

[theo-brown]

Super, can do. Give me a couple of mins

[theo-brown]

Done!