[Feature Request] Add option to use `functorch` for `BasicEnsemble`

[luisenp]

🚀 Feature Request

Change BasicEnsemble to optionally use functorch.vmap.

Motivation and Pitch

Is your feature request related to a problem? Please describe.

BasicEnsemble lets the user provide arbitrary models, which are stacked together using a very naive loop-based implementation. We should be able to do this more efficiently now using functorch.

[natolambert]

Torch has a vmap??? !!!

[luisenp]

It's via functorch which is a separate library but I think they work closely with PyTorch's team.