Add code-wide switch for PyTorch support

[CosmoMatt]

Currently we force users to install both JAX and PyTorch to gain access to either. We should add a switch with optional dependencies to reduce the requirement overhead.

[lgrcia]

Hi @CosmoMatt, is there any plan to work on that in the near future? I would be happy to help if you have any interest in a pull request.

[CosmoMatt]

Hey @lgrcia so this is something we've been wondering about internally, mainly from an implementation point of view (i.e. whether to wrap the JAX backend or just translate it outright). I suspect either approach should be relatively straightforward, but we don't have the time currently to do it. If you are interested in doing this we're of course more than happy for a PR! If you have any specific questions definitely post here and I'll try to follow up.