mlx
mlx copied to clipboard
Indexing an `mlx.array` with list
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!
Indeed.. I see no reason not to support using a list to index.
Sure, great. I will start working on that