jitfields.sym

Overview

This module contains linear-algebra routines (matrix-vector product, matrix inversion, etc.) for batches of symmetric matrices stored in a compact way (that is, only \(N(N+1)/2\) values are stored, instead of \(N^2\)).

Our compact representation differs from classical "columns" or "rows" layouts. The compact flattened matrix should contain the diagonal of the matrix first, followed by the rows of the upper triangle of the matrix, i.e.:

[ a d e ]
[ . b f ]   =>  [a b c d e f]
[ . . c ]

Note that matrix-vector functions (matvec, solve) also accept (and automatically detect) compact diagonal matrices and compact scaled identity matrices. If the vector has shape (*, N) and the matrix has shape (*, NN), where * is any number of leading batch dimensions, NN can take values:

• 1: and the matrix is assumed to be a scaled identity,
• N: and the matrix is assumed to be a diagonal matrix,
• N*(N+1)//2: and the matrix is assumed to be symmetric,
• N*N: and the matrix is assumed to be full.

sym_matvec

sym_matvec(mat, vec, dtype=None, out=None)

Matrix-vector product for compact symmetric matrices

Equivalent to out = mat @ vec.

Parameters:

Name	Type	Description	Default
mat	`(..., CC) tensor`	Symmetric matrix with compact storage, with shape (..., CC), where CC is one of {1, C, C*(C+1)//2, C*C}. If CC == 1, the matrix is a scaled identity. If CC == C, the matrix is diagonal. If CC == C*(C+1)//2, the matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle. If CC == C*C, the matrix is a flattened dense matrix.	required
vec	`(..., C) tensor`	Vector with shape (..., C).	required
dtype	`torch.dtype`	Data type used to carry the computation. By default, same as input.	None
out	`(..., C) tensor`	Output placeholder, with shape (..., C).	None

Returns:

Name	Type	Description
out	(..., C) tensor	Matrix-vector product, with shape (..., C).

sym_addmatvec

sym_addmatvec(inp, mat, vec, dtype=None, out=None)

Add a matrix-vector product for compact symmetric matrices

Equivalent to out = inp + mat @ vec

Parameters:

Name	Type	Description	Default
inp	`(..., C) tensor`	Vector to which the matrix-vector product is added. With shape (..., C).	required
mat	`(..., CC) tensor`	Symmetric matrix with compact storage, with shape (..., CC), where CC is one of {1, C, C*(C+1)//2, C*C}. If CC == 1, the matrix is a scaled identity. If CC == C, the matrix is diagonal. If CC == C*(C+1)//2, the matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle. If CC == C*C, the matrix is a flattened dense matrix.	required
vec	`(..., C) tensor`	Vector with shape (..., C).	required
dtype	`torch.dtype`	Data type used to carry the computation. By default, same as input.	None
out	`(..., C) tensor`	Output placeholder, with shape (..., C).	None

Returns:

Name	Type	Description
out	`(..., C) tensor`	Added matrix-vector product, with shape (..., C).

sym_addmatvec_

sym_addmatvec_(inp, mat, vec, dtype=None)

Inplace add a matrix-vector product for compact symmetric matrices

Equivalent to inp += mat @ vec

Parameters:

Name	Type	Description	Default
inp	`(..., C) tensor`	Vector to which the matrix-vector product is added. With shape (..., C).	required
mat	`(..., CC) tensor`	Symmetric matrix with compact storage, with shape (..., CC), where CC is one of {1, C, C*(C+1)//2, C*C}. If CC == 1, the matrix is a scaled identity. If CC == C, the matrix is diagonal. If CC == C*(C+1)//2, the matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle. If CC == C*C, the matrix is a flattened dense matrix.	required
vec	`(..., C) tensor`	Vector with shape (..., C).	required
dtype	`torch.dtype`	Data type used to carry the computation. By default, same as input.	None

Returns:

Name	Type	Description
inp	`(..., C) tensor`	Added matrix-vector product, with shape (..., C).

sym_submatvec

sym_submatvec(inp, mat, vec, dtype=None, out=None)

Subtract a matrix-vector product for compact symmetric matrices

Equivalent to out = inp - mat @ vec

Parameters:

Name	Type	Description	Default
inp	`(..., C) tensor`	Vector to which the matrix-vector product is added. With shape (..., C).	required
mat	`(..., CC) tensor`	Symmetric matrix with compact storage, with shape (..., CC), where CC is one of {1, C, C*(C+1)//2, C*C}. If CC == 1, the matrix is a scaled identity. If CC == C, the matrix is diagonal. If CC == C*(C+1)//2, the matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle. If CC == C*C, the matrix is a flattened dense matrix.	required
vec	`(..., C) tensor`	Vector with shape (..., C).	required
dtype	`torch.dtype`	Data type used to carry the computation. By default, same as input.	None
out	`(..., C) tensor`	Output placeholder, with shape (..., C).	None

Returns:

Name	Type	Description
out	`(..., C) tensor`	Subtracted matrix-vector product, with shape (..., C).

sym_submatvec_

sym_submatvec_(inp, mat, vec, dtype=None)

Inplace subtract a matrix-vector product for compact symmetric matrices

Equivalent to inp -= mat @ vec

Parameters:

Name	Type	Description	Default
inp	`(..., C) tensor`	Vector to which the matrix-vector product is added. With shape (..., C).	required
mat	`(..., CC) tensor`	Symmetric matrix with compact storage, with shape (..., CC), where CC is one of {1, C, C*(C+1)//2, C*C}. If CC == 1, the matrix is a scaled identity. If CC == C, the matrix is diagonal. If CC == C*(C+1)//2, the matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle. If CC == C*C, the matrix is a flattened dense matrix.	required
vec	`(..., C) tensor`	Vector with shape (..., C).	required
dtype	`torch.dtype`	Data type used to carry the computation. By default, same as input.	None

Returns:

Name	Type	Description
inp	`(..., C) tensor`	Subtracted matrix-vector product, with shape (..., C).

sym_solve

sym_solve(mat, vec, diag=None, dtype=None, out=None)

Solve the symmetric linear system

Equivalent to out = (mat + diag).inverse() @ vec

Warning

Does not backpropagate through mat.

Parameters:

Name	Type	Description	Default
mat	`(..., CC) tensor`	Symmetric matrix with compact storage, with shape (..., CC), where CC is one of {1, C, C*(C+1)//2, C*C}. If CC == 1, the matrix is a scaled identity. If CC == C, the matrix is diagonal. If CC == C*(C+1)//2, the matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle. If CC == C*C, the matrix is a flattened dense matrix.	required
vec	`(..., C) tensor`	Vector with shape (..., C).	required
diag	`(..., C) tensor`	Diagonal regularizer, with shape (..., C).	None
dtype	`torch.dtype`	Data type used to carry the computation. By default, same as input.	None
out	`(..., C) tensor`	Output placeholder, with shape (..., C).	None

Returns:

Name	Type	Description
out	`(..., C) tensor`	Solution of the linear system, with shape (..., C).

sym_solve_

sym_solve_(mat, vec, diag=None, dtype=None)

Solve the symmetric linear system in-place

Equivalent to vec = mat.inverse() @ vec

Warning

Does not backpropagate through mat.

Parameters:

Name	Type	Description	Default
mat	`(..., CC) tensor`	Symmetric matrix with compact storage, with shape (..., CC), where CC is one of {1, C, C*(C+1)//2, C*C}. If CC == 1, the matrix is a scaled identity. If CC == C, the matrix is diagonal. If CC == C*(C+1)//2, the matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle. If CC == C*C, the matrix is a flattened dense matrix.	required
vec	`(..., C) tensor`	Vector with shape (..., C).	required
diag	`(..., C) tensor`	Diagonal regularizer, with shape (..., C).	None
dtype	`torch.dtype`	Data type used to carry the computation. By default, same as input.	None

Returns:

Name	Type	Description
vec	(..., C) tensor	Solution of the linear system, with shape (..., C).

sym_invert

sym_invert(mat, dtype=None, out=None)

Invert a compact symmetric matrix

Equivalent to out = mat.inverse()

Warning

Does not backpropagate through mat.

Parameters:

Name	Type	Description	Default
mat	`(..., C*(C+1)//2) tensor`	Symmetric matrix with compact storage, with shape (..., C*(C+1)//2). The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.	required
dtype	dtype	Data type used to carry the computation. By default, same as input.	None
out	`(..., C*(C+1)//2) tensor`	Output placeholder, with shape (..., C*(C+1)//2).	None

Returns:

Name	Type	Description
mat	`(..., C*(C+1)//2) tensor`	Inverse matrix, with shape (..., C*(C+1)//2).

sym_invert_

sym_invert_(mat, dtype=None)

Invert a compact symmetric matrix in-place

Equivalent to mat = mat.inverse()

Warning

Does not backpropagate through mat.

Parameters:

Name	Type	Description	Default
mat	`(..., C*(C+1)//2) tensor`	Symmetric matrix with compact storage, with shape (..., C*(C+1)//2). The matrix should be saved as a vector containing the diagonal followed by the rows of the upper triangle.	required
dtype	`torch.dtype`	Data type used to carry the computation. By default, same as input.	None

Returns:

Name	Type	Description
mat	`(..., C*(C+1)//2) tensor`	Inverse matrix, with shape (..., C*(C+1)//2).