Fix vmap constant output size

ml-explore / mlx

MLX: An array framework for Apple silicon

https://ml-explore.github.io/mlx/

MIT License

17.59k stars 1.02k forks source link

Closed barronalex closed 4 weeks ago

barronalex commented 1 month ago

Resolves #1516

This mirror's Jax's vmap behaviour so I think it's reasonable.