mlx.core.expand_dims

mlx.core.expand_dims(a: array, /, axis: Union[int, List[int]], *, stream: Union[None, Stream, Device] = None) → array

Add a size one dimension at the given axis.

Parameters:

• a (array) – Input array.

• axes (int or tuple(int)) – The index of the inserted dimensions.

Returns:

The array with inserted dimensions.

Return type:

array