[Feature] out of tree torch models using torchax

[samos123]

Checklist

• 1. If the issue you raised is not a feature but a question, please raise a discussion at https://github.com/sgl-project/sgl-jax/discussions/new/choose Otherwise, it will be closed.
• 2. Please use English, otherwise it will be closed.

Motivation

We have a custom torch model that we wish to run on TPU. the vLLM tpu-inference project uses torchax and similarly torchax may be used for sglang-jax to support torch models.

Why?

• Support a wider range of models out of the box OSS models without having to port them to jax.
• Ability to users to run their private torch models on TPU using sglang-jax.

Related resources

No response