torch_openreml.covariance.Matrix¶

class torch_openreml.covariance.Matrix(shape, param_specs)[source]¶

Bases: ABC

Abstract base class for covariance matrices with parameterized structure.

\[\symbf{V} = \symbf{V}(\symbf{\theta})\]

where \(\symbf{\theta}\) denotes the collection of variance component parameters that define the matrix entries.

This class provides utilities for parameter validation, transform application, and Jacobian computation (both manual and automatic). Subclasses must implement __call__() to construct their specific matrix structure from the provided parameters.

Initialize a covariance matrix with parameter specifications.

Parameters:

shape (tuple or None) – Expected output dimensions of the constructed matrix. Used for validation; the actual shape may be set by subclasses.
param_specs (dict) – Parameter specifications. Keys should be non-empty strings representing parameter names, which must not contain "/". Values should be dictionaries containing the specification for each parameter. Each specification dictionary should contain the keys "fixed", "default", and "trans", representing whether the parameter is fixed or free (bool), the default value (1D torch.Tensor), and the transform (Transform), respectively.

Raises:

TypeError – If param_specs does not follow any of the requirements listed in the argument description, or if shape is not a tuple or torch.Size.
ValueError – If shape values are non-negative, or if a parameter name is empty or contains "/".

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_params`([free_params, include_fixed, ...])	Construct the full parameter tensor from free parameters.
`disable_cache`()	Disable caching of intermediate computation results.
`enable_cache`()	Enable caching of intermediate computation results.
`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.
`inverse`([free_params])	Construct the matrix and return its inverse.
`manual_grad`([free_params])	Compute the Jacobian of `__call__()` with respect to free parameters using a closed-form analytic expression.
`reset_intermediates`()	Clear the intermediate computation cache.
`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.

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.
`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()`.

grad_mode¶

Gradient computation mode. "manual" uses a class-defined manual gradient, "auto" uses automatic differentiation, and "default" uses the manual gradient if manual_grad() is defined, otherwise automatic differentiation.

Type:: str

jacobian_method¶

Jacobian method used by auto_grad(), one of "jacrev", "jacfwd", or "jacobian".

Type:: str

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. Ignored by "jacfwd", which takes no chunk size argument, and by "jacobian", which already loops over the outputs one at a time.

Type:: int or None

cache¶

Whether intermediate computation results are cached, toggled by enable_cache() and disable_cache(). When False, set_intermediates() does nothing and get_intermediates() always returns None, both before validating their arguments.

Type:: bool

set_intermediates(params, intermediates)[source]¶

Cache intermediate computation results keyed by the parameter values.

Stores arbitrary intermediate values alongside the current parameter tensor, its dtype, and its device. Cached values can be retrieved via get_intermediates() to avoid redundant computation across multiple calls with identical parameters.

When cache is False the request is ignored: nothing is stored and params is not validated, so a malformed argument is silently accepted.

Parameters:

params (torch.Tensor) – Current parameter tensor.
intermediates – Arbitrary object to cache (e.g. Cholesky factors, eigendecompositions, or any reusable computation).

Raises:

TypeError – If params is not a Torch tensor and cache is True.
ValueError – If params is not a 1D tensor and cache is True.

Note

If params has length 0 (no free parameters), this is a no-op. A copy of params is stored, so a later in-place edit of the caller’s tensor cannot change the key of an entry already cached.

Example:

import torch
from torch_openreml.covariance import DiagonalMatrix

mat = DiagonalMatrix(3)
free_params = torch.tensor([0.0, 0.5, 1.0])
params = mat.build_params(free_params)
mat.set_intermediates(params, {"log(sigma^2)/2": torch.log(params) / 2})
mat.get_intermediates(params)

{'log(sigma^2)/2': tensor([0.0000, 0.5000, 1.0000])}

get_intermediates(params)[source]¶

Retrieve cached intermediate computation results if still valid.

Compares the values, dtype, and device of params against the stored cache from the last set_intermediates() call. Returns the cached value only if all three match, ensuring stale results are never returned after a parameter update, device transfer, or dtype cast.

Values are compared elementwise, so an entry is reused only when the parameters agree in every position: the same values in a different order, or the same values paired differently, are different models and each gets its own entry. Parameter vectors are a handful of floats, so the comparison is cheap.

When cache is False, None is returned without validating params, so a malformed argument is silently ignored.

Parameters:

params (torch.Tensor) – Current parameter tensor.

Raises:

TypeError – If params is not a Torch tensor and cache is True.
ValueError – If params is not a 1D tensor and cache is True.

Returns:

The cached intermediate object if the cache is valid, or None if the cache is missing, stale, disabled, or params has length 0.

Example:

import torch
from torch_openreml.covariance import DiagonalMatrix

mat = DiagonalMatrix(3)
free_params = torch.tensor([0.0, 0.5, 1.0])
params = mat.build_params(free_params)
mat.set_intermediates(params, {"log(sigma^2)/2": torch.log(params) / 2})
mat.get_intermediates(params)

{'log(sigma^2)/2': tensor([0.0000, 0.5000, 1.0000])}

reset_intermediates()[source]¶

Clear the intermediate computation cache.

Resets all cached values set by set_intermediates() to None, forcing subsequent calls to get_intermediates() to return None until the cache is repopulated. Called automatically in __init__() and within auto_grad() before triggering a fresh Jacobian computation.

Example:

import torch
from torch_openreml.covariance import DiagonalMatrix

mat = DiagonalMatrix(3)
free_params = torch.tensor([0.0, 0.5, 1.0])
params = mat.build_params(free_params)
mat.set_intermediates(params, {"log(sigma^2)/2": torch.log(params) / 2})
print(mat.get_intermediates(params))

{'log(sigma^2)/2': tensor([0.0000, 0.5000, 1.0000])}

mat.reset_intermediates()
print(mat.get_intermediates(free_params))

None

enable_cache()[source]¶

Enable caching of intermediate computation results.

Sets cache to True, so set_intermediates() stores entries again and get_intermediates() can return them. Nothing is cached until the next set_intermediates() call, so entries cleared by an earlier disable_cache() do not come back.

Example:

from torch_openreml.covariance import ScalarMatrix

mat = ScalarMatrix(3)
mat.disable_cache()
mat.enable_cache()
mat.cache

True

disable_cache()[source]¶

Disable caching of intermediate computation results.

Sets cache to False and clears the cache via reset_intermediates(), so an entry stored while caching was enabled cannot be returned if caching is turned back on. While disabled, set_intermediates() does nothing and get_intermediates() always returns None.

Example:

import torch
from torch_openreml.covariance import ScalarMatrix

mat = ScalarMatrix(3)
params = mat.build_params(torch.tensor([0.5]))
mat.set_intermediates(params, "value")
mat.disable_cache()
print(mat.get_intermediates(params))

None

get_default_dtype_device()[source]¶

Get the dtype and the device of the parameter defaults.

All parameter defaults are expected to share the same dtype and device, so that they can be concatenated into a single tensor. When the matrix has no parameters at all, the default dtype and device of Torch are returned instead.

Returns:: (device, dtype) shared by the parameter defaults, or the default device and dtype of Torch when the matrix has no parameters.
Return type:: tuple
Raises:: ValueError – If the parameter defaults do not all share the same dtype and device.

Example:

from torch_openreml.covariance import ScalarMatrix

mat = ScalarMatrix(3)
mat.get_default_dtype_device()

(device(type='cpu'), torch.float32)

build_params(free_params=None, include_fixed=True, trans=True, out_format='tensor')[source]¶

Construct the full parameter tensor from free parameters.

Merges free (trainable) parameters with fixed parameter defaults and applies parameter transforms. Optionally returns a dictionary mapping parameter names to their transformed values.

Parameters:

free_params (torch.Tensor or dict) – Flat 1D tensor of free parameters or a dictionary mapping parameter names to tensors. If omitted, default values are used. Default: None.
include_fixed (bool, optional) – Whether to include fixed parameters in the output. Default: True.
trans (bool, optional) – Whether to apply parameter transforms to the output. Default: True.
out_format (str, optional) – Output format. One of "tensor" or "dict". Default: "tensor".

Returns:

Full parameter tensor of length num_params (or num_free_params when include_fixed=False), or a dictionary mapping parameter names to value tensors.

Return type:

torch.Tensor or dict

Raises:

ValueError – If out_format is not "tensor" or "dict".
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 DiagonalMatrix

mat = DiagonalMatrix(3)
free_params = torch.tensor([0.0, 0.5, 1.0])
mat.build_params(free_params)

tensor([1.0000, 2.7183, 7.3891])

mat.build_params()

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

mat.param_specs["sigma^2_2"]["fixed"] = True
mat.build_params(free_params[0:2])

tensor([1.0000, 2.7183, 1.0000])

mat.build_params(free_params[0:2], include_fixed=False)

tensor([1.0000, 2.7183])

mat.build_params(free_params[0:2], include_fixed=False, trans=False)

tensor([0.0000, 0.5000])

set_param_specs(key, *, fixed=None, default=None, trans=None)[source]¶

Update the fields of one or more parameter specifications in place.

key is resolved against param_specs, and each field passed as something other than None replaces that field of every resolved specification. Fields left as None are kept unchanged, so a call without fields changes nothing. Values are assigned as given, without validation or copying, and take effect on the next call.

A key that does not end with "/" addresses one specification exactly. A key that ends with "/" addresses a subtree instead: every parameter whose key has key as a leading path, including the node itself. "A/" therefore addresses every parameter of the operand "A", nested operands below it included, and "/" addresses every parameter of the matrix.

Parameters:

key (str) – Key of the parameter specification, as param_specs exposes it, or a subtree path ending with "/". Composites namespace their operands’ parameters, so "A/sigma^2" addresses the parameter "sigma^2" of the operand "A", and "A/" addresses every parameter under it.
fixed (bool, optional) – New value for the "fixed" field. Default: None.
default (torch.Tensor, optional) – New value for the "default" field. Default: None.
trans (Transform, optional) – New value for the "trans" field. Default: None.

Raises:

TypeError – If key is not a str.
ValueError – If key is not a parameter of this matrix, and not a subtree path that resolves to at least one.

Example:

import torch
from torch_openreml.covariance import DiagonalMatrix

mat = DiagonalMatrix(3)
mat.set_param_specs("sigma^2_1", fixed=True, default=torch.tensor([7.5]))
mat.free_param_names

['sigma^2_0', 'sigma^2_2']

mat.build_params(torch.tensor([1.0, 3.0]), trans=False)

tensor([1.0000, 7.5000, 3.0000])

from torch_openreml.covariance import BlockDiagonal, ScalarMatrix

op = BlockDiagonal(subject=DiagonalMatrix(2), noise=ScalarMatrix(2))
op.set_param_specs("subject/", fixed=True)
op.free_param_names

['noise/sigma^2']

op.set_param_specs("/", fixed=True)
op.free_param_names

[]

trans_grad(free_params=None)[source]¶

Compute the element-wise derivative of the free parameter transforms.

Evaluates the derivative of each free parameter’s transform function at the current parameter values. Used in the chain rule when computing manual gradients of the matrix with respect to the original (untransformed) parameterisation.

Parameters:

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

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.

Returns:

1D tensor of element-wise transform derivatives, of the same length as free_params. An empty 1D tensor is returned when the matrix has no free parameters.

Return type:

torch.Tensor

Example:

import torch
from torch_openreml.covariance import DiagonalMatrix

mat = DiagonalMatrix(3)
free_params = torch.tensor([0.0, 0.5, 1.0])
mat.trans_grad(free_params)

tensor([ 2.0000,  5.4366, 14.7781])

mat.trans_grad()

tensor([2., 2., 2.])

auto_grad(free_params=None)[source]¶

Compute the Jacobian of build() with respect to free parameters using automatic differentiation.

Uses jacobian_method and jacobian_chunk_size to compute the full Jacobian.

If the matrix has no free parameters, returns (None, []).

Parameters:

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

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.

Returns:

(grad, grad_names), where grad is a 3D tensor of shape (num_free_params, *shape), and grad_names has the same length as grad. (None, []) is returned when the matrix has no free parameters.

Return type:

tuple

Example:

import torch
from torch_openreml.covariance import DiagonalMatrix

mat = DiagonalMatrix(2)
free_params = torch.tensor([0.0, 0.5])
grad, grad_names = mat.auto_grad(free_params)
grad, grad_names

(tensor([[[2.0000, 0.0000],
          [0.0000, 0.0000]],
 
         [[0.0000, 0.0000],
          [0.0000, 5.4366]]]),
 ['sigma^2_0', 'sigma^2_1'])

manual_grad(free_params=None)[source]¶

Compute the Jacobian of __call__() with respect to free 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:

Validate free_params via build_params() first, so bad input raises even when the matrix has no free parameters.
Return (None, []) when the matrix has no free parameters.
Return a 3D gradient tensor of shape (num_free_params, *shape) and a matching list of parameter names.
Apply transform derivatives from trans_grad() via the chain rule so that gradients are with respect to the raw (untransformed) parameters.

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, []) when the matrix has no free parameters.
Return type:: tuple
Raises:: NotImplementedError – If the subclass does not provide an analytic gradient. grad() catches this and falls back to auto_grad().

abstractmethod __call__(free_params=None)[source]¶

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

grad(free_params=None)[source]¶

Compute the Jacobian of __call__() with respect to trainable parameters.

Dispatches to manual_grad() or auto_grad() according to grad_mode:

"default": attempts manual_grad(), falling back to auto_grad() if not implemented.
"auto": always uses auto_grad().

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) as described in manual_grad() and auto_grad().
Return type:: tuple
Raises:: RuntimeError – If grad_mode is not a recognised value.

Example:

import torch
from torch_openreml.covariance import DiagonalMatrix

mat = DiagonalMatrix(2)
free_params = torch.tensor([0.0, 0.5])
grad, grad_names = mat.grad(free_params)
grad, grad_names

(tensor([[[2.0000, 0.0000],
          [0.0000, 0.0000]],
 
         [[0.0000, 0.0000],
          [0.0000, 5.4366]]]),
 ['sigma^2_0', 'sigma^2_1'])

inverse(free_params=None)[source]¶

Construct the matrix and return its inverse.

This is the default implementation: it builds \(\symbf{V}\) via __call__() and inverts the result directly. Subclasses should override it with a cheaper analytic inverse where one is available.

Parameters:

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

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.

Returns:

Inverse of the constructed matrix, of shape shape.

Return type:

torch.Tensor

resolve_dim(n)[source]¶

Resolve a matrix dimension from an integer or a vector of levels.

The resolved dimension is the number of distinct elements of n (see n_distinct()) when n is a list, a tuple or a pandas.Series, and n itself otherwise.

Parameters:: n (int, list, tuple or pandas.Series) – Matrix dimension, or a vector of levels whose distinct elements define it.
Returns:: The resolved matrix dimension.
Return type:: int

Example:

from torch_openreml.covariance import ScalarMatrix

mat = ScalarMatrix(3)
mat.resolve_dim(["a", "b", "a", "c"])

property shape¶

Output matrix shape.

Type:: tuple

property param_specs¶

Parameter specifications, returned as stored rather than copied.

A specification edited in place (for example setting "fixed" to True) therefore takes effect on the next call.

Type:: dict

property param_names¶

Parameter names.

Type:: list of str

property free_param_names¶

Free parameter names.

Type:: list of str

property fixed_param_names¶

Fixed parameter names.

Type:: list of str

property free_param_index¶

Index of free parameters.

Type:: tuple

property fixed_param_index¶

Index of fixed parameters.

Type:: tuple

property num_params¶

Total number of parameters.

Type:: int

property num_free_params¶

Total number of free parameters.

Type:: int

property num_fixed_params¶

Total number of fixed parameters.

Type:: int

property param_defaults¶

Parameter defaults.

Type:: Dict of torch.Tensor

property free_param_defaults¶

Free parameter defaults.

Type:: Dict of torch.Tensor

property fixed_param_defaults¶

Fixed parameter defaults.

Type:: Dict of torch.Tensor

property param_trans¶

Parameter transforms.

Type:: Dict of Transform

property free_param_trans¶

Transforms for free parameters.

Type:: Dict of Transform

property fixed_param_trans¶

Transforms for fixed parameters.

Type:: Dict of Transform

property repr_dict¶

Key-value pairs used to build the string representation.

Type:: dict