mlx.core.take

take(a: array, /, indices: int | array, axis: int | None = None, *, stream: None | Stream | Device = None) → array

沿轴取出元素。

元素沿指定的轴从 indices 中取出。如果未指定轴，则在执行取出操作之前，将数组视为扁平的一维数组。

例如，如果 axis=1，则这等同于 a[:, indices, ...]。

参数：

• a (array) – 输入数组。

• indices (int 或 array) – 整型索引或具有整型类型的输入数组。

• axis (int, 可选) – 执行取出操作的轴。如果未指定，则将数组视为扁平的一维向量。

返回值：

数组 a 的索引值。

返回类型：

array