[Feature Request] Parallel training using vmap

[Benjamin-eecs]

Motivation

speedup solution for training small neural networks on one GPU #17

Solution

• functorch + torchopt

Resource

• functorch example
• functorch zou example
• JAX + FLAX

Checklist

• I have checked that there is no similar issue in the repo (required)