torch_openreml.covariance.LinearPropagation

class torch_openreml.covariance.LinearPropagation(operands)

Bases: Operator

Initialize a covariance matrix with optional parameter transforms.

Parameters:

• shape (tuple or None) – Expected output dimensions of the constructed matrix. Used for validation; the actual shape may be set by subclasses.

• param_names (list of str) – Ordered names of parameters in params. Empty list if no trainable parameters (e.g., fixed matrices).

• trans (list of Transform or None) – List of transforms applied to each parameter before constructing the matrix. If None, no transforms are used. Typically used for variance (\(\exp(2\theta) > 0\)) or correlation constraints (\(\rho \in (-1, 1)\)).

• no_grad_index (list of int) – Indices to exclude from gradient computation. Parameters at these indices will be omitted from grad and grad_names. Use set_no_grad() instead for convenience.

Note

The transform applies as

\[\symbf{V} = \left[f_0(\theta_0), \ldots, f_{p-1}(\theta_{p-1}) \right]^\top,\]

where each \(f_i\) is the i-th transform in trans. If trans has length 1, the single transform is broadcast and applied elementwise to all parameters.

Raises:

• TypeError – If param_names is not a list of strings, or if transforms contain non-Transform objects.

• ValueError – If parameter names are not unique, or if indices in no_grad_index are out of range.

Methods

__call__(params)	Construct the matrix from a flat parameter tensor.
auto_grad(params)	Compute the Jacobian of build() with respect to trainable parameters using automatic differentiation.
build_operands(params)
check_operands(operands)
check_params(params)	Validate a parameter tensor and return its device and dtype.
from_param_dict(param_dict)	Extract parameter tensors from a dictionary into a flat 1D tensor.
get_intermediates(params)	Retrieve cached intermediate computation results if still valid.
grad(params)	Compute the Jacobian of __call__() with respect to trainable parameters.
manual_grad(params)	Compute the Jacobian of __call__() with respect to trainable parameters using a closed-form analytic expression.
map_theta_to_dv(theta)	An interface compatible with torch_openreml.REML that maps parameters to the matrix Jacobian.
map_theta_to_v(theta)	An interface compatible with torch_openreml.REML that maps parameters to a matrix.
operands_grad(params)
reset_intermediates()	Clear the intermediate computation cache.
set_intermediates(params, intermediates)	Cache intermediate computation results keyed by parameter hash.
set_no_grad([index, param_name])	Set the indices of parameters to exclude from gradient computation.
to_param_dict(params)	Convert a flat parameter tensor to a parameter dictionary.
trans_grad(params)	Compute the element-wise derivative of the parameter transforms.
trans_params(params)	Apply parameter transforms to a flat parameter tensor.

Attributes

no_grad_index	Indices of parameters excluded from gradient computation.
num_params	Total number of parameters.
operands
param_names	Ordered parameter names.
repr_dict	Key-value pairs used to build the string representation.
shape	Output matrix shape.
trans	Parameter transforms.

__call__(params)

Construct the matrix from a flat parameter tensor.

Must be implemented by subclasses. Implementations should convert params via from_param_dict() or to_param_dict(), then call check_params() to validate and trans_params() to apply transforms before any computation.

Parameters:

params (torch.Tensor or dict) – Flat 1D parameter tensor or parameter dictionary.

Returns:

Constructed matrix of shape shape.

Return type:

torch.Tensor

manual_grad(params)

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

This method is optional. When implemented by a subclass, grad() will invoke it in preference to auto_grad() under the default grad mode. If not implemented, calling this method raises NotImplementedError and grad() falls back to automatic differentiation.

Implementations must satisfy the following contract:

• Return (None, []) if all parameters are excluded via no_grad_index.

• Return a 3D gradient tensor of shape (num_params - len(no_grad_index), *shape) and a matching list of parameter names, omitting any index in no_grad_index.

• Apply transform derivatives from trans_grad() via the chain rule so that gradients are with respect to the raw (untransformed) parameters.

Parameters:

params (torch.Tensor or dict) – Flat 1D parameter tensor or parameter dictionary.

Returns:

(grad, grad_names), where grad is a 3D tensor of shape (num_params - len(no_grad_index), *shape) and grad_names is a list of the corresponding parameter names. Returns (None, []) if all parameters are excluded from gradient computation.

Return type:

tuple

Raises:

NotImplementedError – If the subclass does not provide an analytic gradient. grad() catches this and falls back to auto_grad().