mlx icon indicating copy to clipboard operation
mlx copied to clipboard

meshgrid and open grid functionality?

Open Datamance opened this issue 2 years ago • 2 comments

First of all - this library is a dream come true for a grad student. I've been using Taichi for doing things on the GPU up to this point (which works well for a lot of things - it's a fantastic library) but having vectorized/tensor computations optimized for apple hardware is going to be a game-changer.

Now, for my question: I see that the mlx.core.array constructor can take a numpy.ndarray, which allows us to do lots of things in numpy first and then convert over to mlx arrays. That said - especially for huge grids with fine granularity - it would be nice to have meshgrid, ogrid, etc. natively implemented. Are there currently plans to implement this (and other) array/tensor conveniences from the numpy and torch apis?

Thanks again!

Datamance avatar Dec 12 '23 20:12 Datamance

We definitely are working towards including more things in MLX natively - especially common numpy ops. Things like linspace, meshgrid etc. We currently have our hands pretty full, but I will include this as enhancement and hopefully we can get to it soon. Also, PRs for any of those are more than welcome!

awni avatar Dec 12 '23 20:12 awni