Adding types to vmap decorator/adapter

[mishmish66]

Very nice tool! Thank you so much for building this.

This is an initial idea for a way to make the vmapping functionality work better with typecheckers.

Likely needs tweaks but I figured I'd put it here so we can discuss it, and I created a discussion thread #31

[mishmish66]

Since decorators cannot add kwargs as per PEP 612 it's not possible to have the adapt_with_vmap decorator support manual dim specifications.

A potential workaround is to accept a positional dict with dict[str, int] type appended to the end of the args with param size specifications?