'Serverless" jax?

[JesseLu]

Hi!

I've been using jax for hpc purposes (solving Maxwell's equations) and it's been amazing. As I'm trying to scale up what I've built so far, I keep on coming back to the idea of being able to do "serverless" jax.

That is, that a user would lower a jax computation to stableHLO, send it along with input data to a GPU equipped server, and receive back the output arrays. Although this wouldn't work well for ML purposes where there would be too much model data, it really seems to fit well for some hpc workloads.

So finally my question: what am I missing and what have I overlooked? And is this the right approach to the problem?