torch_openreml.covariance.SimpleMatrix¶

class torch_openreml.covariance.SimpleMatrix(num_free_params, call, manual_grad=None, default=None)[source]¶

Bases: Matrix

A covariance matrix for simple, function-based parameterisations.

This is the easiest way to use MarginalREML with a custom covariance structure: provide the number of parameters and a function that maps a flat parameter tensor to the covariance matrix. All parameters are free and use an identity transform (unconstrained). The default argument sets the value used for each free parameter when none are provided.

For more advanced needs (custom transforms, fixed parameters, manual gradients), subclass Matrix directly.

Initialize a simple covariance matrix.

Parameters:

num_free_params (int) – Number of free parameters.
call (callable) – Function with signature call(free_params) -> torch.Tensor that constructs the covariance matrix from a flat 1D parameter tensor. It receives the output of build_params() with include_fixed=False and trans=False, so a parameter dict or None is resolved to the free parameter tensor itself, untransformed, before the call.
manual_grad (callable, optional) – Function with signature manual_grad(free_params) -> (grad, grad_names) for a closed-form Jacobian. If None (default), automatic differentiation is used.
default (torch.Tensor, optional) – Default value for each parameter, given as a 1D tensor of shape (1,). Passed to simple_param_specs(). Defaults to torch.tensor([0.0]).

Example:

import torch
from torch_openreml.covariance import SimpleMatrix

def my_v(free_params):
    return torch.diag(free_params)

mat = SimpleMatrix(num_free_params=3, call=my_v, default=torch.tensor([1.0]))
mat(torch.tensor([1.0, 2.0, 3.0]))

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

mat()

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

mat({
    "theta_0": torch.tensor([2.0]),
    "theta_1": torch.tensor([3.0]),
    "theta_2": torch.tensor([4.0]),
})

tensor([[2., 0., 0.],
        [0., 3., 0.],
        [0., 0., 4.]])

mat.grad(torch.tensor([1.0, 2.0, 3.0]))

(tensor([[[1., 0., 0.],
          [0., 0., 0.],
          [0., 0., 0.]],
 
         [[0., 0., 0.],
          [0., 1., 0.],
          [0., 0., 0.]],
 
         [[0., 0., 0.],
          [0., 0., 0.],
          [0., 0., 1.]]]),
 ['theta_0', 'theta_1', 'theta_2'])

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

__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

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