Utilities¶

Scalar maps applied to eigenvalues (with their derivatives, used by the Daleckii–Krein backwards), projections and retractions on the Stiefel manifold (used by the BiMap parametrizations), and random generators for tests and synthetic data.

Scalar functions¶

Scalar functions applied element-wise to eigenvalues: sqrt, inv_sqrt, softplus, and derivatives.

sqrt_derivative(x)[source]¶

Derivative of the square root function.

\[\frac{d}{dx}\sqrt{x} = \frac{1}{2\sqrt{x}}\]

Parameters:: x (torch.Tensor) – Scalar or array of scalars.
Returns:: x_sqrt_deriv (torch.Tensor) – Derivative of sqrt of x
Return type:: Tensor

inv_sqrt(x)[source]¶

Inverse of the square root.

\[f(x) = \frac{1}{\sqrt{x}} = x^{-1/2}\]

Parameters:: x (torch.Tensor) – Scalar or array of scalars.
Returns:: x_inv_sqrt (torch.Tensor) – Inverse sqrt of x
Return type:: Tensor

inv_sqrt_derivative(x)[source]¶

Derivative of the inverse of the square root.

\[\frac{d}{dx} x^{-1/2} = -\frac{1}{2} x^{-3/2}\]

Parameters:: x (torch.Tensor) – Scalar or array of scalars.
Returns:: x_inv_sqrt_deriv (torch.Tensor) – Derivative of the inverse sqrt of x
Return type:: Tensor

inv(x)[source]¶

Inverse function.

\[f(x) = \frac{1}{x}\]

Parameters:: x (torch.Tensor) – Scalar or array of scalars.
Returns:: x_inv (torch.Tensor) – Inverse of x
Return type:: Tensor

scaled_softplus(x)[source]¶

Scaled SoftPlus function.

\[f(x) = \log_2\big(1 + 2^{x}\big)\]

Base-2 (rather than the usual base-\(e\)) SoftPlus, chosen so that \(f(0) = 1\), \(f(x) \to 0\) as \(x \to -\infty\), and \(f'(x) \to 1\) as \(x \to +\infty\). Used (via scaled_softplus_symmetric()) to reparametrize eigenvalues so BiMap weights stay strictly positive.

Parameters:: x (torch.Tensor) – Scalar or array of scalars.
Returns:: x_softplus (torch.Tensor) – SoftPlus of x
Return type:: Tensor

scaled_softplus_derivative(x)[source]¶

Derivative of the scaled SoftPlus function.

\[f'(x) = \sigma(x \ln 2)\]

where \(\sigma\) is the logistic sigmoid.

Parameters:: x (torch.Tensor) – Scalar or array of scalars.
Returns:: x_softplus_deriv (torch.Tensor) – Derivative of SoftPlus of x
Return type:: Tensor

inv_scaled_softplus(x)[source]¶

Inverse of the scaled SoftPlus function.

\[f^{-1}(x) = \log_2\big(2^{x} - 1\big)\]

Inverse of scaled_softplus().

Parameters:: x (torch.Tensor) – Scalar or array of scalars.
Returns:: x_inv_softplus (torch.Tensor) – Inverse of SoftPlus of x
Return type:: Tensor

inv_scaled_softplus_derivative(x)[source]¶

Derivative of the inverse of the scaled SoftPlus function.

\[(f^{-1})'(x) = \frac{1}{1 - 2^{-x}}\]

Parameters:: x (torch.Tensor) – Scalar or array of scalars.
Returns:: x_softplus_deriv (torch.Tensor) – Derivative of the inverse SoftPlus of x
Return type:: Tensor

Stiefel manifold¶

autograd path	manual backward
`stiefel_projection_polar()`	`StiefelProjectionPolar`	Projection from the ambient space onto the Stiefel manifold based on the polar decomposition
`stiefel_projection_tangent_orthogonal()`	`StiefelProjectionTangentOrthogonal`	Orthogonal projection from the ambient space onto the tangent space of the Stiefel manifold at point
`stiefel_projection_qr()`	`StiefelProjectionQR`	Projection from the ambient space onto the Stiefel manifold based on the QR decomposition
`stiefel_differential_projection_qr()`	–	Differential of the projection on Stiefel based on QR decomposition
`stiefel_adjoint_differential_projection_qr()`	–	Adjoint of the differential projection on Stiefel based on QR decomposition

Projections and retractions on the Stiefel manifold (orthonormal matrices).

stiefel_projection_polar(point)[source]¶

Projection from the ambient space onto the Stiefel manifold based on the polar decomposition

Parameters:: point (torch.Tensor of shape (n_in, n_out)) – Rectangular matrix (with n_out <= n_in)
Returns:: projected_point (torch.Tensor of shape (n_in, n_out)) – Orthogonal matrix
Return type:: Tensor

stiefel_projection_tangent_orthogonal(vector, point)[source]¶

Orthogonal projection from the ambient space onto the tangent space of the Stiefel manifold at point

Note that this also corresponds to both the differential and differential adjoint of projection_stiefel_polar

Parameters:

vector (torch.Tensor of shape (n_in, n_out)) – Rectangular matrix (direction)
point (torch.Tensor of shape (n_in, n_out)) – Orthogonal matrix (with n_out <= n_in)

Returns:

tangent_vector (torch.Tensor of shape (n_in, n_out)) – Rectangular matrix

Return type:

Tensor

class StiefelProjectionTangentOrthogonal(*args, **kwargs)[source]¶

Orthogonal projection from the ambient space onto the tangent space of the Stiefel manifold at point

static forward(ctx, vector, point)[source]¶

Forward pass of the orthogonal projection from the ambient space onto the tangent space of the Stiefel manifold at point

Parameters:

ctx (torch.autograd.function._ContextMethodMixin) – Context object to retrieve tensors saved during the forward pass
vector (torch.Tensor of shape (n_in, n_out)) – Rectangular matrix (direction)
point (torch.Tensor of shape (n_in, n_out)) – Orthogonal matrix (with n_out <= n_in)

Returns:

tangent_vector (torch.Tensor of shape (n_in, n_out)) – Rectangular matrix

Return type:

Tensor

static backward(ctx, grad_output)[source]¶

Backward pass of the orthogonal projection onto the tangent space of the Stiefel manifold at point

Parameters:

ctx (torch.autograd.function._ContextMethodMixin) – Context object to retrieve tensors saved during the forward pass
grad_output (torch.Tensor of shape (n_in, n_out)) – Gradient of the loss with respect to the output tangent vector

Returns:

grad_input (torch.Tensor of shape (n_in, n_out)) – Gradient of the loss with respect to the input rectangular matrix

Return type:

tuple[Tensor, None]

class StiefelProjectionPolar(*args, **kwargs)[source]¶

Projection from the ambient space onto the Stiefel manifold based on polar decomposition

static forward(ctx, point)[source]¶

Forward pass of the projection onto the Stiefel manifold

Parameters:

ctx (torch.autograd.function._ContextMethodMixin) – Context object to retrieve tensors saved during the forward pass
point (torch.Tensor of shape (n_in, n_out)) – Rectangular matrix (with n_out <= n_in)

Returns:

projected_point (torch.Tensor of shape (n_in, n_out)) – Orthogonal matrix

Return type:

Tensor

static backward(ctx, grad_output)[source]¶

Backward pass of the projection onto the Stiefel manifold

Parameters:

ctx (torch.autograd.function._ContextMethodMixin) – Context object to retrieve tensors saved during the forward pass
grad_output (torch.Tensor of shape (n_in, n_out)) – Gradient of the loss with respect to the projected orthogonal matrix

Returns:

grad_input (torch.Tensor of shape (n_in, n_out)) – Gradient of the loss with respect to the input rectangular matrix

Return type:

Tensor

stiefel_projection_qr(point)[source]¶

Projection from the ambient space onto the Stiefel manifold based on the QR decomposition

Parameters:: point (torch.Tensor of shape (n_in, n_out)) – Rectangular matrix (with n_out <= n_in)
Returns:: projected_point (torch.Tensor of shape (n_in, n_out)) – Orthogonal matrix
Return type:: Tensor

stiefel_differential_projection_qr(vector, Q, R)[source]¶

Differential of the projection on Stiefel based on QR decomposition

Parameters:

vector (torch.Tensor of shape (n_in, n_out)) – Rectangular matrix (direction)
Q (torch.Tensor of shape (n_in, n_out)) – Q factor of RQ decomposition of point (with n_out <= n_in)
R (torch.Tensor of shape (n_out, n_out)) – R factor of QR decomposition of point (upper triangular)

Returns:

tangent_vector (torch.Tensor of shape (n_in, n_out)) – Rectangular matrix

Return type:

Tensor

stiefel_adjoint_differential_projection_qr(vector, Q, R)[source]¶

Adjoint of the differential projection on Stiefel based on QR decomposition

Parameters:

vector (torch.Tensor of shape (n_in, n_out)) – Rectangular matrix (direction)
Q (torch.Tensor of shape (n_in, n_out)) – Q factor of RQ decomposition of point (with n_out <= n_in)
R (torch.Tensor of shape (n_out, n_out)) – R factor of QR decomposition of point (upper triangular)

Returns:

transformed_vector (torch.Tensor of shape (n_in, n_out)) – Rectangular matrix

Return type:

Tensor

class StiefelProjectionQR(*args, **kwargs)[source]¶

Projection from the ambient space onto the Stiefel manifold based on QR decomposition

static forward(ctx, point)[source]¶

Forward pass of the projection onto the Stiefel manifold

Parameters:

ctx (torch.autograd.function._ContextMethodMixin) – Context object to retrieve tensors saved during the forward pass
point (torch.Tensor of shape (n_in, n_out)) – Rectangular matrix (with n_out <= n_in)

Returns:

projected_point (torch.Tensor of shape (n_in, n_out)) – Orthogonal matrix

Return type:

Tensor

static backward(ctx, grad_output)[source]¶

Backward pass of the projection onto the Stiefel manifold

Parameters:

ctx (torch.autograd.function._ContextMethodMixin) – Context object to retrieve tensors saved during the forward pass
grad_output (torch.Tensor of shape (n_in, n_out)) – Gradient of the loss with respect to the projected orthogonal matrix

Returns:

grad_input (torch.Tensor of shape (n_in, n_out)) – Gradient of the loss with respect to the input rectangular matrix

Return type:

Tensor

Random matrices¶

Random SPD matrices with a prescribed condition number.

random_DPD(n_features, n_matrices=1, cond=10, device=None, dtype=None, generator=None)[source]¶

Generate a batch of random diagonal positive definite matrices

Parameters:

n_features (int) – Number of features
n_matrices (int, optional) – Number of matrices. Default is 1
cond (float, optional) – Condition number w.r.t. inversion of SPD matrices. Default is 10
device (torch.device, optional) – Torch device. Default is None.
dtype (torch.dtype, optional) – Torch dtype. Default is None.
generator (torch.Generator, optional) – Generator to ensure reproducibility. Default is None.

Returns:

torch.Tensor of shape (n_matrices, n_features, n_features) – Batch of DPD matrices

Return type:

Tensor

random_SPD(n_features, n_matrices=1, cond=10, device=None, dtype=None, generator=None)[source]¶

Generate a batch of random SPD matrices

Parameters:

n_features (int) – Number of features
n_matrices (int, optional) – Number of matrices. Default is 1
cond (float, optional) – Condition number w.r.t. inversion of SPD matrices. Default is 10
device (torch.device, optional) – Torch device. Default is None.
dtype (torch.dtype, optional) – Torch dtype. Default is None.
generator (torch.Generator, optional) – Generator to ensure reproducibility. Default is None.

Returns:

torch.Tensor of shape (n_matrices, n_features, n_features) – Batch of SPD matrices

Return type:

Tensor

Random points on the Stiefel manifold and Stiefel weight initialization.

random_stiefel(n_in, n_out, n_matrices=1, device=None, dtype=None, generator=None)[source]¶

Random point on Stiefel manifold

Parameters:

n_in (int) – Number of rows
n_out (int) – Number of columns
n_matrices (int, optional) – Number of matrices. Default is 1
device (torch.device, optional) – Torch device. Default is None
dtype (torch.dtype, optional) – Torch dtype. Default is None
generator (torch.Generator, optional) – Generator to ensure reproducibility. Default is None

Return type:

Tensor