nitorch_fastmath.sugar

Overview

This module contains "syntactic sugar" linear algebra routines. I.e., their behaviour is relatively easily to implement in pure PyTorch, but sometimes not very readable as it involves squeezing, unsqueezing, or shuffling dimensions.

For example, we implement matvec(mat, vec), which in PyTorch must be written matmul(mat, vec.unsqueeze(-1)).squeeze(-1).

kron2

kron2(a, b)

Kronecker product of two matrices \(\mathbf{A} \otimes \mathbf{B}\)

Parameters:

Name	Type	Description	Default
`a`	`Tensor`	Left matrix, with shape `(..., m, n)`.	required
`b`	`(..., p, q) tensor`	Right matrix, with shape `(..., p, q)`.	required

Returns:

Name	Type	Description
`ab`	`(..., pm, qn) tensor`	Kronecker product, with shape `(..., pm, qn)`. Note: `ab.reshape([P, M, Q, N])[p, m, q, n] == a[m, n] * b[n, q]`

References

Kronecker product, Wikipedia.

lmdiv

lmdiv(a, b, method='lu', rcond=1e-15, out=None)

Left matrix division \(\mathbf{A}^{-1} \times \mathbf{B}\).

Note

if m != n, the Moore-Penrose pseudoinverse is always used.

Parameters:

Name	Type	Description	Default
`a`	`(..., m, n) tensor`	Left input ("the system"), with shape `(..., m, n)`.	required
`b`	`(..., m, k) tensor`	Right input ("the point"), with shape `(..., m, k)`.	required
`method`	`{'lu', 'chol', 'svd', 'pinv'}`	Inversion method: `'lu'` : LU decomposition. `a` must be invertible. `'chol'` : Cholesky decomposition. `a` must be positive definite. `'svd'` : Singular Value decomposition. `'pinv'` : Moore-Penrose pseudoinverse (by means of SVD with thresholded singular values).	`'lu'`
`rcond`	`float`	Cutoff for small singular values when `method='pinv'`.	`1e-15`
`out`	`tensor`	Output tensor.	`None`

Returns:

Name	Type	Description
`x`	`(..., n, k) tensor`	Solution of the linear system, with shape `(..., n, k)`,

References

LU decomposition, Wikipedia.
Cholesky decomposition, Wikipedia.
Singular value decomposition, Wikipedia.
Moore-Penrose inverse, Wikipedia.

rmdiv

rmdiv(a, b, method='lu', rcond=1e-15, out=None)

Right matrix division \(\mathbf{A} \times \mathbf{B}^{-1}\).

Note

if m != n, the Moore-Penrose pseudoinverse is always used.

Parameters:

Name	Type	Description	Default
`a`	`(..., k, m) tensor`	Left input ("the point"), with shape `(..., k, m)`.	required
`b`	`(..., n, m) tensor`	Right input ("the system"), with shape `(..., n, m)`.	required
`method`	`{'lu', 'chol', 'svd', 'pinv'}`	Inversion method: `'lu'` : LU decomposition. `a` must be invertible. `'chol'` : Cholesky decomposition. `a` must be positive definite. `'svd'` : Singular Value decomposition. `'pinv'` : Moore-Penrose pseudoinverse (by means of SVD with thresholded singular values).	`'lu'`
`rcond`	`float`	Cutoff for small singular values when `method='pinv'`.	`1e-15`
`out`	`tensor`	Output tensor	`None`

Returns:

Name	Type	Description
`x`	`(..., k, m) tensor`	Solution of the linear system, with shape `(..., k, m)`.

References

LU decomposition, Wikipedia.
Cholesky decomposition, Wikipedia.
Singular value decomposition, Wikipedia.
Moore-Penrose inverse, Wikipedia.

inv

inv(a, method='lu', rcond=1e-15, out=None)

Matrix inversion \(\mathbf{A}^{-1}\).

Note

if m != n, the Moore-Penrose pseudoinverse is always used.

Parameters:

Name	Type	Description	Default
`a`	`(..., m, n) tensor`	Input matrix, with shape `(..., m, n)`.	required
`method`	`{'lu', 'chol', 'svd', 'pinv'}`	Inversion method: `'lu'` : LU decomposition. `a` must be invertible. `'chol'` : Cholesky decomposition. `a` must be positive definite. `'svd'` : Singular Value decomposition. `'pinv'` : Moore-Penrose pseudoinverse (by means of SVD with thresholded singular values).	`'lu'`
`rcond`	`float`	Cutoff for small singular values when `method='pinv'`.	`1e-15`
`out`	`tensor`	Output tensor).	`None`

Returns:

Name	Type	Description
`x`	`(..., n, m) tensor`	Inverse matrix, with shape `(..., n, m)`.

References

LU decomposition, Wikipedia.
Cholesky decomposition, Wikipedia.
Singular value decomposition, Wikipedia.
Moore-Penrose inverse, Wikipedia.

matvec

matvec(mat, vec, out=None)

Matrix-vector product \(\mathbf{A} \times \mathbf{b}\) (supports broadcasting) .

Parameters:

Name	Type	Description	Default
`mat`	`(..., m, n) tensor`	Input matrix, with shape `(..., m, n)`.	required
`vec`	`(..., n) tensor`	Input vector, with shape `(..., n)`.	required
`out`	`(..., n) tensor`	Placeholder for the output tensor, with shape `(..., n)`.	`None`

Returns:

Name	Type	Description
`mv`	`(..., m) tensor`	Matrix vector product of the inputs, with shape `(..., m)`.

solvevec

solvevec(mat, vec, method='lu', rcond=1e-15, out=None)

Left matrix-vector division \(\mathbf{A}^{-1} \times \mathbf{b}\).

Note

if m != n, the Moore-Penrose pseudoinverse is always used.

Parameters:

Name	Type	Description	Default
`mat`	`(..., m, n) tensor`	Left input ("the system"), with shape `(..., m, n)`.	required
`vec`	`(..., m) tensor`	Right input ("the point"), with shape `(..., m)`.	required
`method`	`{'lu', 'chol', 'svd', 'pinv'}`	Inversion method: `'lu'` : LU decomposition. `a` must be invertible. `'chol'` : Cholesky decomposition. `a` must be positive definite. `'svd'` : Singular Value decomposition. `'pinv'` : Moore-Penrose pseudoinverse (by means of SVD with thresholded singular values).	`'lu'`
`rcond`	`float`	Cutoff for small singular values when `method='pinv'`.	`1e-15`
`out`	`tensor`	Output tensor.	`None`

Returns:

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

References

LU decomposition, Wikipedia.
Cholesky decomposition, Wikipedia.
Singular value decomposition, Wikipedia.
Moore-Penrose inverse, Wikipedia.

outer

outer(a, b, out=None)

Outer product of two (batched) tensors \(\mathbf{a} \otimes \mathbf{b} = \mathbf{a} \times \mathbf{b}^\mathrm{H}\).

Warning

When \(\mathbf{a}\) and \(\mathbf{b}\) are complex, this function returns the complex outer product \(\mathbf{a} \otimes \mathbf{b} = \mathbf{a}\mathbf{b}^\mathrm{H}\).

Parameters:

Name	Type	Description	Default
`a`	`(..., n) tensor`	Left vector, with shape `(..., n)`	required
`b`	`(..., m) tensor`	Right vector, with shape `(..., m)`	required
`out`	`(..., n, m) tensor`	Output placeholder, with shape `(..., n, m)`	`None`

Returns:

Name	Type	Description
`out`	`(..., n, m) tensor`	Outer product, with shape `(..., n, m)`

References

Outer product, Wikipedia.

trace

trace(a, keepdim=False)

Compute the (batched) trace of a matrix \(\operatorname{tr}\left(\mathbf{A}\right)\).

Parameters:

Name	Type	Description	Default
`a`	`(..., m, m) tensor`		required

Returns:

Name	Type	Description
`t`	`(..., [1, 1]) tensor`

References

Trace, Wikipedia.

dot

dot(a, b, keepdim=False, out=None)

(Batched) dot product of two vectors \(\mathbf{a}^\mathrm{H}\mathbf{b}\).

Warning

When \(\mathbf{a}\) and \(\mathbf{b}\) are complex, this function returns the complex dot product \((\mathbf{a}, \mathbf{b}) = \mathbf{a}^\mathrm{H}\mathbf{b}\). This dot product is linear in the second term and antilinear in the first term.

This behaviour differs from torch.dot and torch.inner, which both compute \(\mathbf{a}^\mathrm{T}\mathbf{b}\) when the vectors are complex.

Parameters:

Name	Type	Description	Default
`a`	`(..., n) tensor`	Left vector, with shape `(..., n)`	required
`b`	`(..., m) tensor`	Right vector, with shape `(..., m)`	required
`keepdim`	`bool`	Keep reduced dimension	`False`
`out`	`tensor`	Output placeholder	`None`

Returns:

Name	Type	Description
`ab`	`(..., [1]) tensor`	Dot product, with shape `(..., [1])`.

References

Dot product, Wikipedia.

mdot

mdot(a, b, keepdim=False, out=None)

Compute the Frobenius inner product of two matrices \(\operatorname{tr}(\mathbf{A}^\mathrm{H}\mathbf{B})\).

Warning

When \(\mathbf{A}\) and \(\mathbf{B}\) are complex, this function returns the complex dot product \((\mathbf{A}, \mathbf{B}) = \operatorname{tr}(\mathbf{A}^\mathrm{H}\mathbf{B})\). This dot product is linear in the second term and antilinear in the first term.

Parameters:

Name	Type	Description	Default
`a`	`(..., n, m) tensor`	Left matrix, with shape `(..., n, m)`.	required
`b`	`(..., n, m) tensor`	Right matrix, with shape `(..., n, m)`.	required
`keepdim`	`bool`	Keep reduced dimensions.	`False`
`out`	`(..., [1, 1]) tensor`	Output placeholder.	`None`

Returns:

Name	Type	Description
`dot`	`(..., [1, 1]) tensor`	Matrix inner product, with shape `(..., [1, 1])`.

References

Frobenius inner product, Wikipedia.

is_orthonormal

is_orthonormal(basis, return_matrix=False)

Check that a basis is an orthonormal basis.

Parameters:

Name	Type	Description	Default
`basis`	`(F, N, [M]) tensor`	A basis of a vector or matrix space, with shape `(F, N, [M])`. `F` is the number of elements in the basis.	required
`return_matrix`	`bool`	If True, return the matrix of all pairs of inner products between elements if the basis	`False`

Returns:

Name	Type	Description
`check`	`bool`	True if the basis is orthonormal
`matrix`	`(F, F) tensor`, if `return_matrix is True`	Matrix of all pairs of inner products, with shape `(F, F) tensor`.

round

round(t, decimals=0)

Round a tensor to the given number of decimals.

Parameters:

Name	Type	Description	Default
`t`	`tensor`	Input tensor.	required
`decimals`	`int`	Round to this decimal.	`0`

Returns:

Name	Type	Description
`t`	`tensor`	Rounded tensor.