search

ml-explore / mlx

MLX: An array framework for Apple silicon
https://ml-explore.github.io/mlx/
MIT License
17.37k stars 1.01k forks source link

Fix concatenate vmap #1600

Closed angeloskath closed 2 days ago

angeloskath commented 3 days ago

Fixes #1599. I do think the concatenate vmap was wrong to not expand/broadcast the non-vmapped inputs. I added an example test that fails before this PR similar to the example in #1599.