torch_openreml.covariance.Transpose

class torch_openreml.covariance.Transpose(*args, **kwargs)

Bases: Operator

Transpose of a covariance matrix.

\[\symbf{V} = \symbf{A}^\top\]

If \(\symbf{A}\) is \(m \times n\), the result is an \(n \times m\) matrix. The operand may be a trainable Matrix instance or a fixed torch.Tensor value.

Initialize a transpose operator from exactly one operand.

Parameters:

• *args – Exactly one operand as a positional argument or a single dict. The operand is \(\symbf{A}\).

• **kwargs – Exactly one operand as a keyword argument.

Raises:

ValueError – If the number of operands is not exactly one.

Example:

import torch
from torch_openreml.covariance import DummyMatrix, Transpose

op = Transpose(a=DummyMatrix(["a", "b", "c", "a"]))
op(torch.tensor([]))

tensor([[1., 0., 0., 1.],
        [0., 1., 0., 0.],
        [0., 0., 1., 0.]])

Methods

__call__([free_params])	Construct the matrix from a flat parameter tensor.
auto_grad([free_params])	Compute the Jacobian of build() with respect to free parameters using automatic differentiation.
build_operands([free_params])	Evaluate each operand at the current free parameters.
build_params([free_params, include_fixed, ...])	Construct the full parameter tensor by delegating to each operand.
disable_cache()	Disable intermediate caching on this operator and its operands.
enable_cache()	Enable intermediate caching on this operator and its operands.
get_default_dtype_device()	Get the dtype and the device of the parameter defaults.
get_intermediates(params)	Retrieve cached intermediate computation results if still valid.
grad([free_params])	Compute the Jacobian of __call__() with respect to trainable parameters.
grad_tree([free_params])	Compute the gradient of this operator and of every nested operand, returning each keyed by its path.
inverse([free_params])	Return the inverse of the composite matrix from its closed form.
manual_grad([free_params])	Compute the Jacobian of __call__() with respect to trainable parameters using a closed-form analytic expression.
operands_grad([free_params])	Compute the Jacobian of each operand with respect to its parameters.
operands_inverse([free_params])	Compute the inverse of each operand at the current free parameters.
reset_intermediates()	Clear the intermediate caches of this operator and its operands.
resolve_dim(n)	Resolve a matrix dimension from an integer or a vector of levels.
set_intermediates(params, intermediates)	Cache intermediate computation results keyed by the parameter values.
set_param_specs(key, *[, fixed, default, trans])	Update the fields of one or more parameter specifications in place.
trans_grad([free_params])	Compute the element-wise derivative of the free parameter transforms.
tree([free_params])	Evaluate this operator and every nested operand, returning each result keyed by its path.

Attributes

fixed_param_defaults	Fixed parameter defaults.
fixed_param_index	Index of fixed parameters.
fixed_param_names	Fixed parameter names.
fixed_param_trans	Transforms for fixed parameters.
free_param_defaults	Free parameter defaults.
free_param_index	Index of free parameters.
free_param_names	Free parameter names.
free_param_trans	Transforms for free parameters.
num_fixed_params	Total number of fixed parameters.
num_free_params	Total number of free parameters.
num_params	Total number of parameters.
operands	Mapping from operand names to operand matrices or tensors.
param_defaults	Parameter defaults.
param_names	Parameter names.
param_specs	Parameter specifications, returned as stored rather than copied.
param_trans	Parameter transforms.
repr_dict	Key-value pairs used to build the string representation.
shape	Output matrix shape.
grad_mode	Gradient computation mode.
jacobian_method	Jacobian method used by auto_grad(), one of "jacrev", "jacfwd", or "jacobian".
jacobian_chunk_size	Number of outputs held per batch when auto_grad() differentiates with "jacrev", or None to compute the Jacobian in a single batch.
cache	Whether intermediate computation results are cached, toggled by enable_cache() and disable_cache().

__call__(free_params=None)

Construct the matrix from a flat parameter tensor.

Must be implemented by subclasses. Implementations should convert free_params via build_params() to validate, include fixed parameters, and apply transforms before any computation.

Parameters:

free_params (torch.Tensor or dict) – Flat 1D parameter tensor or parameter dictionary. If omitted, default values are used. Default: None.

Returns:

Constructed matrix of shape shape.

Return type:

torch.Tensor

inverse(free_params=None)

Return the inverse of the composite matrix from its closed form.

The composite is \(\symbf{V} = \symbf{A}^\top\), so its inverse is \(\symbf{V}^{-1} = (\symbf{A}^{-1})^\top\): the operand is inverted through its own inverse() and the result is transposed, so the composite is never inverted densely.

Warning

The operand must be square, otherwise the composite has no inverse and raises.

Parameters:

free_params (torch.Tensor or dict) – Flat 1D parameter tensor or parameter dictionary. If omitted, default values are used. Default: None.

Returns:

The n x n inverse of the composite matrix.

Return type:

torch.Tensor

Example:

import torch
from torch_openreml.covariance import AR1Matrix, Transpose

op = Transpose(a=AR1Matrix(3))
free_params = torch.tensor([0.0, 0.5])
op.inverse(free_params)

tensor([[ 1.0638, -0.2605,  0.0000],
        [-0.2605,  1.1276, -0.2605],
        [ 0.0000, -0.2605,  1.0638]])

manual_grad(free_params=None)

Compute the Jacobian of __call__() with respect to trainable parameters using a closed-form analytic expression.

If \(\symbf{V} = \symbf{A}^\top\), the gradient with respect to \(\theta_{\symbf{A}}\) is the transpose of the operand’s own Jacobian:

\[\frac{\partial \symbf{V}}{\partial \theta_{\symbf{A}}} = \left(\frac{\partial \symbf{A}}{\partial \theta_{\symbf{A}}}\right)^\top\]

The per-operand Jacobian from operands_grad() is transposed over its last two dimensions.

Parameters:

free_params (torch.Tensor or dict) – Flat 1D parameter tensor or parameter dictionary. If omitted, default values are used. Default: None.

Returns:

(grad, grad_names), where grad is a 3D tensor of shape (num_free_params, *shape) and grad_names is a list of the corresponding parameter names. Returns (None, []) if the operand has no free parameters.

Return type:

tuple

Raises:

• TypeError – If free_params is not a Torch tensor.

• ValueError – If free_params is not a 1D tensor or has the wrong length, or if free_params is a dict with missing or unexpected keys.

Example:

import torch
from torch_openreml.covariance import LowerTriangularMatrix, Transpose

op = Transpose(a=LowerTriangularMatrix(3, 2))
free_params = torch.tensor([0.0, 0.5, 1.0, 0.2, -0.3])
grad, grad_names = op.manual_grad(free_params)
grad

tensor([[[1., 0., 0.],
         [0., 0., 0.]],

        [[0., 1., 0.],
         [0., 0., 0.]],

        [[0., 0., 0.],
         [0., 1., 0.]],

        [[0., 0., 1.],
         [0., 0., 0.]],

        [[0., 0., 0.],
         [0., 0., 1.]]])

grad_names

['a/L_0_0', 'a/L_1_0', 'a/L_1_1', 'a/L_2_0', 'a/L_2_1']