Function vmap_trace 

pub fn vmap_trace(
    function: &MlxClosure,
    inputs: &MlxArrayVector,
    input_axes: &[i32],
) -> Result<(MlxArrayVector, MlxArrayVector, MlxVectorInt), MlxCError>

Traces a closure for the VMAP transform, producing traced inputs and outputs plus the resolved axes.

Errors

Returns the captured MLX-C description when tracing fails.