spd_learn.modules.Vec#

class spd_learn.modules.Vec(device: device | None = None, dtype: dtype | None = None)[source]#

Bases: Module

Vectorization Layer.

This layer vectorizes a batch of matrices along the last two dimensions.

forward(X: Tensor) → Tensor[source]#

Forward pass of the Vec layer.

Parameters:

X (torch.Tensor) – A batch of matrices with shape (…, n, k).

Returns:

A batch of vectorized matrices with shape (…, n * k).

Return type:

torch.Tensor

inverse_transform(X: Tensor, n_rows: int) → Tensor[source]#

Inverse transform of the Vec layer.

Parameters:

  • X (torch.Tensor) – A batch of vectorized matrices with shape (…, n_rows * k).

  • n_rows (int) – The number of rows in the original matrices.

Returns:

A batch of matrices with shape (…, n_rows, k).

Return type:

torch.Tensor