mlx icon indicating copy to clipboard operation
mlx copied to clipboard

Indexing an `mlx.array` with list

Open amirhossein-razlighi opened this issue 1 year ago • 2 comments

If we index an mlx.core.array type with a list, an exception occurs. While common convention in numpy or PyTorch is to return the given elements correctly. E.g:

import mlx.core as mx
import numpy as np
import torch

a = mlx.array([1,2,3])
a[[0,1]] # ValueError: Cannot index mlx array using the given type.

b  = np.array([1,2,3])
b[[0,1]] # array([1, 2])

c = torch.tensor([1,2,3])
c[[0,1]] # tensor([1, 2])

If it's a valid issue and we want to follow these convention in mlx, please let me know. I'll be happy to work on it!

amirhossein-razlighi avatar Mar 26 '24 20:03 amirhossein-razlighi

Indeed.. I see no reason not to support using a list to index.

awni avatar Mar 26 '24 20:03 awni

Sure, great. I will start working on that

amirhossein-razlighi avatar Mar 26 '24 20:03 amirhossein-razlighi