Rebuild vmap for Rune to support nesting, JIT, and grad

[tmattio]

Summary

• Replace the current WIP vmap in rune/lib/vmap.ml with a clean implementation that matches the behaviour of Jax's vmap.
• Ensure the new design composes with higher-order transforms such as jit, grad, jvp, and nested vmap.

Problem

• The existing code in rune/lib/vmap.ml is mostly AI-generated and fails the nested scenarios in rune/test/test_vmap.ml.
• I suspect it currently does not compose well with jit or grad, preventing us from stacking transformations.

Proposed Scope

• Investigate design space for vmap and propose an effect-based solution that's compatible with Jax semantics and compose well with our other effect handlers.
• Implement the design in vmap.ml and expose through Rune.vmap. Make sure the implementation passes the tests at test_vmap.ml

[akshayamadhuri]

@tmattio i would like to work on it

[tmattio]

@akshayamadhuri before that, can you confirm on Discord that you've managed to build and run the tests (you can share a screenshot), once you've confirmed I'll be able to assign you to issues