Support torch.vmap

[zihaozou]

Currently, GSplat supports single-scene, multi-camera rendering in parallel. However, in some cases, we need to render multiple scenes with multiple cameras in parallel, such as for deformable scenes. A simple approach is to use a for loop, but a better option is to use torch.vmap, which enables parallelism for unbatched operations.

However, to use vmap with a custom torch.autograd.Function, we need to make a small adjustment. As specified here, we simply need to override an additional static method, setup_context, in the custom Function class.

If you're interested, I can help create a PR.

[maturk]

I would be interested in seeing how to do this, or a small demo :)