Making torch dep optional?

[dlwh]

Hi,

For those of us JAX cultists, any chance we could make torch an optional dependency? I think the only place the main framework needs torch is for the distributed evaluation stuff in evaluator.py, which can be import-guarded with a few lines' fix.

[dlwh]

(I'd be happy to PR this if you all are interested)

[lintangsutawika]

Happy to review your PR. Also it would be great if there is some sort of example on how you using in jax for beginners, just an idea.

[haileyschoelkopf]

Hey @dlwh ! Echoing this--happy to see a PR making this change, and we can discuss there whether the changes are self-contained enough to merge generally. Should be a good idea to support JAX.

[dlwh]

@lintangsutawika Sounds like a good idea. I'm gonna PR the core change first just as a POC and once I have things working on my end I'll try to make a tutorial