Add tests and benchmarks for TorchBench models

[zou3519]

Inspired by https://github.com/facebookresearch/torchdynamo/blob/main/torchbench.py

There should be some way to:

1. Iterate over the TorchBench models
2. Run correctness tests and performance comparisons for a number of different tasks (like in the autograd.functional benchmark suite (https://github.com/pytorch/pytorch/blob/master/benchmarks/functional_autograd_benchmark/functional_autograd_benchmark.py).
3. Ideally we run correctness tests somehow as a test in our CI

The various tasks we'd run would be something like:

• grad (basic functorch test)
• vmap x grad (per-sample-grad)
• vmap (ensembling)
• vmap x jacrev (batched jacobian of the input)
• vmap x jacrev (per-sample-jacobian of the parameters)

Things to consider:

• A lot of the models will likely OOM -- how would we handle that?
• The TorchBench suite is representative of NN models that people run, but a lot of functorch use cases are different from "computing interesting quantities on traditional NN models".