torch_openreml.covariance.Sum

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

Bases: Operator

Sum of multiple covariance matrices.

\[\symbf{V} = \sum_{i=1}^{k} \symbf{A}_i\]

where each \(\symbf{A}_i\) is a covariance matrix of the same shape. This operator represents additive covariance structures, commonly used in linear mixed-effects models to combine multiple variance components (e.g., genetic, environmental, and residual).

All operands must evaluate to matrices of identical shape. Each operand may be a trainable Matrix or a fixed torch.Tensor.

Initialize a sum operator from two or more operands.

Parameters:

• *args – Two or more operands as positional arguments or a single dict mapping names to operands.

• **kwargs – Two or more operands as keyword arguments.

Raises:

ValueError – If fewer than two operands are provided.

Example:

import torch
from torch_openreml.covariance import AR1Matrix, ScalarMatrix, Sum

op = Sum(time=AR1Matrix(4), noise=ScalarMatrix(4))
free_params = torch.tensor([0.5, 1.0, 1.0])
op(free_params)

tensor([[10.1073,  1.2562,  0.5805,  0.2683],
        [ 1.2562, 10.1073,  1.2562,  0.5805],
        [ 0.5805,  1.2562, 10.1073,  1.2562],
        [ 0.2683,  0.5805,  1.2562, 10.1073]])

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])	Construct the matrix and return its inverse.
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

manual_grad(free_params=None)

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

Since \(\symbf{V} = \sum_i \symbf{A}_i\), the gradient with respect to each operand’s parameters is simply the operand’s own gradient — there is no cross-term interaction. Per-operand Jacobians from operands_grad() are concatenated directly.

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 all parameters are fixed.

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 AR1Matrix, ScalarMatrix, Sum

op = Sum(time=AR1Matrix(4), noise=ScalarMatrix(4))
free_params = torch.tensor([0.5, 1.0, 1.0])
grad, grad_names = op.manual_grad(free_params)
grad

tensor([[[ 5.4366,  2.5123,  1.1610,  0.5365],
         [ 2.5123,  5.4366,  2.5123,  1.1610],
         [ 1.1610,  2.5123,  5.4366,  2.5123],
         [ 0.5365,  1.1610,  2.5123,  5.4366]],

        [[ 0.0000,  1.0689,  0.9879,  0.6848],
         [ 1.0689,  0.0000,  1.0689,  0.9879],
         [ 0.9879,  1.0689,  0.0000,  1.0689],
         [ 0.6848,  0.9879,  1.0689,  0.0000]],

        [[14.7781,  0.0000,  0.0000,  0.0000],
         [ 0.0000, 14.7781,  0.0000,  0.0000],
         [ 0.0000,  0.0000, 14.7781,  0.0000],
         [ 0.0000,  0.0000,  0.0000, 14.7781]]])

grad_names

['time/sigma^2', 'time/rho', 'noise/sigma^2']