Add initial support for PyTorch backend

[HarshvirSandhu]

Description

Add PyTorch support for few basic Pytensor Ops using torch.compile method and following the existing JAX architecture.
Also replicated this PyTorch sandbox example with this backend.

cc @ricardoV94 @aseyboldt

Type of change

• New feature / enhancement

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[aseyboldt]

Hey :-)
Great to have some code already!
I left a few comments while skimming through the code, just ignore if they are not helpful right now.

[ricardoV94]

First approach looks okay although some things that were copied from JAX don't necessarily make sense for Pytorch. We should remove those.

Also may be good to start adding some tests

[ricardoV94]

This is getting closer to ready. I suggest removing all dispatch that are not being tested as we can add the functionality in follow up PRs

[ricardoV94]

@HarshvirSandhu I'm sorry I didn't mean that we should parametrize all tests with device, only those that are explicitly about inputs/ outputs, which are the first 3 of test_basic:
test_pytorch_FunctionGraph_once
test_shared
test_shared_updates

All others you don't have to worry about device

[ricardoV94]

Looks great!

[ricardoV94]

Some failing tests: AssertionError: https://github.com/pymc-devs/pytensor/actions/runs/9547073282/job/26311326610?pr=764#step:6:287

Torch not compiled with CUDA enabled

I actually don't know if github CI/CD gives us access to GPUs, so we may need to skip the GPU parametrization when the device is not available (so it can still be tested locally if someone has GPU support, but does not fail on Github)

[ricardoV94]

Yeah GPUs are only available on team/enterprise github: https://docs.github.com/en/enterprise-cloud@latest/actions/using-github-hosted-runners/about-larger-runners/about-larger-runners

We need to skip the gpu parametrizations when cuda is not available. You should be able to use pytest.skip inside the test. This is also useful so people can test locally even if they don't have CUDA

[ricardoV94]

Only want to skip on the GPU parametrization

Suggested change

	if torch.cuda.is_available() is False:
	if device == "cuda" and not torch.cuda.is_available():

Same thing in the other tests

[ricardoV94]

Nice work @HarshvirSandhu !