mlx.core.identity

mlx.core.identity(n: int, dtype: Optional[Dtype] = float32, *, stream: Union[None, Stream, Device] = None) → array

Create a square identity matrix.

Parameters:

• n (int) – The number of rows and columns in the output.

• dtype (Dtype, optional) – Data type of the output array. Defaults to float32.

• stream (Stream, optional) – Stream or device. Defaults to None.

Returns:

An identity matrix of size n x n.

Return type:

array