Skip to main content

​
Overview

The mlx.core.linalg module provides a comprehensive set of linear algebra operations including matrix decompositions, norms, solvers, and eigenvalue computations.

​
Matrix Norms

​
norm

mx.linalg.norm(
    a: array,
    ord: float | str = None,
    axis: int | tuple = None,
    keepdims: bool = False,
    stream: StreamOrDevice = None
) -> array
Compute vector or matrix norms.
​
a
array
required
Input array
​
ord
float | str
default:"None"
Order of the norm. For vectors: can be any real number, 'inf', or '-inf'. For matrices: can be 'fro' (Frobenius), 'nuc' (nuclear), 1, -1, 2, -2, 'inf', or '-inf'. If None, computes the 2-norm (or Frobenius norm for matrices).
​
axis
int | tuple
default:"None"
Axis or axes along which to compute the norm. If None and ord is also None, computes norm of flattened array.
​
keepdims
bool
default:"False"
If True, the reduced axes are left in the result as dimensions with size one
​
stream
StreamOrDevice
default:"None"
Stream or device to use
Returns: Norm of the array Example: 
import mlx.core as mx

# Vector 2-norm
x = mx.array([3.0, 4.0])
norm = mx.linalg.norm(x)  # 5.0

# Matrix Frobenius norm
A = mx.array([[1.0, 2.0], [3.0, 4.0]])
norm = mx.linalg.norm(A, ord='fro')

# Infinity norm
norm = mx.linalg.norm(x, ord=float('inf'))  # 4.0

​
Matrix Decompositions

​
qr

mx.linalg.qr(
    a: array,
    stream: StreamOrDevice = None
) -> tuple[array, array]
Compute the QR decomposition of a matrix. Decomposes matrix a into Q @ R where Q is orthogonal and R is upper triangular.
​
a
array
required
Input matrix (must be at least 2D and floating point)
​
stream
StreamOrDevice
default:"None"
Stream or device to use
Returns: Tuple (Q, R) where Q is orthogonal and R is upper triangular Example: 
A = mx.array([[2.0, 3.0], [1.0, 2.0]])
Q, R = mx.linalg.qr(A, stream=mx.cpu)

# Verify decomposition
assert mx.allclose(Q @ R, A)

# Q is orthogonal
assert mx.allclose(Q.T @ Q, mx.eye(2))

​
svd

mx.linalg.svd(
    a: array,
    compute_uv: bool = True,
    stream: StreamOrDevice = None
) -> list[array]
Compute the Singular Value Decomposition (SVD). Decomposes matrix a into U @ diag(S) @ Vt.
​
a
array
required
Input matrix
​
compute_uv
bool
default:"True"
If True, returns [U, S, Vt]. If False, returns only [S].
​
stream
StreamOrDevice
default:"None"
Stream or device to use
Returns: List containing [U, S, Vt] if compute_uv=True, otherwise [S] Example: 
A = mx.array([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]], dtype=mx.float32)
U, S, Vt = mx.linalg.svd(A, stream=mx.cpu)

# Reconstruct matrix
A_reconstructed = U[:, :len(S)] @ mx.diag(S) @ Vt
assert mx.allclose(A_reconstructed, A, rtol=1e-5)

# Just singular values
S = mx.linalg.svd(A, compute_uv=False, stream=mx.cpu)

​
cholesky

mx.linalg.cholesky(
    a: array,
    upper: bool = False,
    stream: StreamOrDevice = None
) -> array
Compute the Cholesky decomposition of a positive definite matrix.
​
a
array
required
Positive definite input matrix
​
upper
bool
default:"False"
If True, compute upper triangular decomposition. Otherwise, compute lower triangular.
​
stream
StreamOrDevice
default:"None"
Stream or device to use
Returns: Lower (or upper) triangular matrix L such that A = L @ L.T (or A = U.T @ U) Example: 
A = mx.array([[4.0, 2.0], [2.0, 3.0]])
L = mx.linalg.cholesky(A)

# Verify decomposition
assert mx.allclose(L @ L.T, A)

​
lu

mx.linalg.lu(
    a: array,
    stream: StreamOrDevice = None
) -> list[array]
Compute the LU decomposition with partial pivoting.
​
a
array
required
Input matrix
​
stream
StreamOrDevice
default:"None"
Stream or device to use
Returns: List [P, L, U] where P is the permutation matrix, L is lower triangular, and U is upper triangular Example: 
A = mx.array([[2.0, 1.0], [1.0, 2.0]])
P, L, U = mx.linalg.lu(A, stream=mx.cpu)

# Verify decomposition
assert mx.allclose(P @ L @ U, A)

​
lu_factor

mx.linalg.lu_factor(
    a: array,
    stream: StreamOrDevice = None
) -> tuple[array, array]
Compute the LU factorization with partial pivoting in compact form.
​
a
array
required
Input matrix
​
stream
StreamOrDevice
default:"None"
Stream or device to use
Returns: Tuple (LU, pivots) where LU contains both L and U in a single matrix, and pivots contains pivot indices Example: 
A = mx.array([[2.0, 1.0], [1.0, 2.0]])
LU, pivots = mx.linalg.lu_factor(A, stream=mx.cpu)

​
Eigenvalues and Eigenvectors

​
eig

mx.linalg.eig(
    a: array,
    stream: StreamOrDevice = None
) -> tuple[array, array]
Compute eigenvalues and eigenvectors of a square matrix.
​
a
array
required
Square input matrix
​
stream
StreamOrDevice
default:"None"
Stream or device to use
Returns: Tuple (eigenvalues, eigenvectors) Example: 
A = mx.array([[1.0, 2.0], [2.0, 1.0]])
eigenvalues, eigenvectors = mx.linalg.eig(A, stream=mx.cpu)

​
eigvals

mx.linalg.eigvals(
    a: array,
    stream: StreamOrDevice = None
) -> array
Compute eigenvalues of a square matrix (without eigenvectors).
​
a
array
required
Square input matrix
​
stream
StreamOrDevice
default:"None"
Stream or device to use
Returns: Array of eigenvalues Example: 
A = mx.array([[1.0, 2.0], [2.0, 1.0]])
eigenvalues = mx.linalg.eigvals(A, stream=mx.cpu)

​
eigh

mx.linalg.eigh(
    a: array,
    UPLO: str = 'L',
    stream: StreamOrDevice = None
) -> tuple[array, array]
Compute eigenvalues and eigenvectors of a Hermitian or symmetric matrix.
​
a
array
required
Hermitian or symmetric input matrix
​
UPLO
str
default:"'L'"
Specifies whether to use the upper (‘U’) or lower (‘L’) triangular part of the matrix
​
stream
StreamOrDevice
default:"None"
Stream or device to use
Returns: Tuple (eigenvalues, eigenvectors) with real eigenvalues Example: 
A = mx.array([[1.0, 2.0], [2.0, 1.0]])
eigenvalues, eigenvectors = mx.linalg.eigh(A, stream=mx.cpu)

​
eigvalsh

mx.linalg.eigvalsh(
    a: array,
    UPLO: str = 'L',
    stream: StreamOrDevice = None
) -> array
Compute eigenvalues of a Hermitian or symmetric matrix (without eigenvectors).
​
a
array
required
Hermitian or symmetric input matrix
​
UPLO
str
default:"'L'"
Specifies whether to use the upper (‘U’) or lower (‘L’) triangular part of the matrix
​
stream
StreamOrDevice
default:"None"
Stream or device to use
Returns: Array of real eigenvalues Example: 
A = mx.array([[1.0, 2.0], [2.0, 1.0]])
eigenvalues = mx.linalg.eigvalsh(A, stream=mx.cpu)

​
Matrix Inversion

​
inv

mx.linalg.inv(
    a: array,
    stream: StreamOrDevice = None
) -> array
Compute the inverse of a square matrix.
​
a
array
required
Square invertible matrix
​
stream
StreamOrDevice
default:"None"
Stream or device to use
Returns: Inverse of the input matrix Example: 
A = mx.array([[1.0, 2.0], [3.0, 4.0]])
A_inv = mx.linalg.inv(A, stream=mx.cpu)

# Verify inversion
assert mx.allclose(A @ A_inv, mx.eye(2))

​
pinv

mx.linalg.pinv(
    a: array,
    stream: StreamOrDevice = None
) -> array
Compute the Moore-Penrose pseudoinverse of a matrix.
​
a
array
required
Input matrix (can be non-square)
​
stream
StreamOrDevice
default:"None"
Stream or device to use
Returns: Pseudoinverse of the input matrix Example: 
A = mx.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
A_pinv = mx.linalg.pinv(A, stream=mx.cpu)

​
tri_inv

mx.linalg.tri_inv(
    a: array,
    upper: bool = False,
    stream: StreamOrDevice = None
) -> array
Compute the inverse of a triangular matrix.
​
a
array
required
Triangular input matrix
​
upper
bool
default:"False"
If True, a is upper triangular. Otherwise, a is lower triangular.
​
stream
StreamOrDevice
default:"None"
Stream or device to use
Returns: Inverse of the triangular matrix Example: 
L = mx.array([[1.0, 0.0], [2.0, 3.0]])
L_inv = mx.linalg.tri_inv(L, upper=False, stream=mx.cpu)

​
cholesky_inv

mx.linalg.cholesky_inv(
    a: array,
    upper: bool = False,
    stream: StreamOrDevice = None
) -> array
Compute the inverse of a positive definite matrix using its Cholesky decomposition.
​
a
array
required
Cholesky factor (lower or upper triangular)
​
upper
bool
default:"False"
If True, a is the upper triangular Cholesky factor. Otherwise, a is lower triangular.
​
stream
StreamOrDevice
default:"None"
Stream or device to use
Returns: Inverse of the original positive definite matrix Example: 
A = mx.array([[4.0, 2.0], [2.0, 3.0]])
L = mx.linalg.cholesky(A)
A_inv = mx.linalg.cholesky_inv(L, stream=mx.cpu)

​
Linear System Solvers

​
solve

mx.linalg.solve(
    a: array,
    b: array,
    stream: StreamOrDevice = None
) -> array
Solve the linear system A @ x = b for x.
​
a
array
required
Coefficient matrix (square)
​
b
array
required
Right-hand side vector or matrix
​
stream
StreamOrDevice
default:"None"
Stream or device to use
Returns: Solution x to the linear system Example: 
A = mx.array([[3.0, 1.0], [1.0, 2.0]])
b = mx.array([9.0, 8.0])
x = mx.linalg.solve(A, b, stream=mx.cpu)

# Verify solution
assert mx.allclose(A @ x, b)

​
solve_triangular

mx.linalg.solve_triangular(
    a: array,
    b: array,
    upper: bool = False,
    stream: StreamOrDevice = None
) -> array
Solve the triangular linear system A @ x = b for x.
​
a
array
required
Triangular coefficient matrix
​
b
array
required
Right-hand side vector or matrix
​
upper
bool
default:"False"
If True, a is upper triangular. Otherwise, a is lower triangular.
​
stream
StreamOrDevice
default:"None"
Stream or device to use
Returns: Solution x to the triangular system Example: 
L = mx.array([[1.0, 0.0], [2.0, 3.0]])
b = mx.array([1.0, 8.0])
x = mx.linalg.solve_triangular(L, b, upper=False, stream=mx.cpu)

​
Other Operations

​
cross

mx.linalg.cross(
    a: array,
    b: array,
    axis: int = -1,
    stream: StreamOrDevice = None
) -> array
Compute the cross product of two arrays.
​
a
array
required
First input array
​
b
array
required
Second input array
​
axis
int
default:"-1"
Axis along which to compute the cross product
​
stream
StreamOrDevice
default:"None"
Stream or device to use
Returns: Cross product of a and b Example: 
a = mx.array([1.0, 0.0, 0.0])
b = mx.array([0.0, 1.0, 0.0])
c = mx.linalg.cross(a, b)  # [0, 0, 1]