Add a qjit-compatible catalyst.vmap function

[maliasadi]

• catalyst.vmap(fun) is added to Catalyst and can apply inside qjitted functions backed by catalyst.for_loop
• catalyst.vmap dispatches to jax.vmap when is called outsize QJIT
• support in_axes
• support out_axes
• support axis_size
• Support PyTrees as parameters and return values

[sc-55115]

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.54%. Comparing base (d191ede) to head (6a8a215).

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #497   +/-   ##
=======================================
  Coverage   99.54%   99.54%           
=======================================
  Files          51       51           
  Lines        8493     8571   +78     
  Branches      572      598   +26     
=======================================
+ Hits         8454     8532   +78     
  Misses         21       21           
  Partials       18       18           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[dime10]

Nice work 🥇 The main function is really well-structured, although due to the complexity of the vmap options some blocks are still a bit unclear to me.

[dime10]

Thanks @maliasadi, this is looking great! I identified a few more issues with the error checking and matching JAX's behaviour, but otherwise the functionality is really nice 🎉

[josh146]

Nice work @maliasadi 😎

[dime10]

Great work 💯