torch_openreml.covariance.Adapter¶

class torch_openreml.covariance.Adapter(adaptee, param_specs, param_map)[source]¶

Bases: Matrix

Reparameterise an existing covariance matrix with a new set of parameters.

An Adapter wraps an adaptee matrix and exposes a user-defined set of parameters. A callable param_map translates the adapter’s parameter tensor into the parameter tensor expected by the adaptee:

\[\symbf{V}(\boldsymbol{\theta}) = \symbf{V}_{\text{adaptee}}\!\big(f(\boldsymbol{\theta})\big)\]

where \(f\) is param_map and \(\boldsymbol{\theta}\) are the adapter’s free parameters.

Gradients are computed via the chain rule: the adaptee’s gradient with respect to its own parameters is pulled back through the Jacobian of param_map.

Initialise an adapter wrapping an existing matrix.

Parameters:

adaptee (Matrix) – The matrix to reparameterise.
param_specs (dict) – Parameter specifications for the adapter’s own parameters. Keys are parameter names; values are dictionaries with keys "fixed", "default", and "trans". All parameters should be free and use TransformIdentity.
param_map (callable) – A function f(params) -> adaptee_params that maps the adapter’s parameter tensor to the parameter tensor expected by adaptee. Must be differentiable (compatible with jacobian_method).

Raises:

ValueError – If any parameter specification uses a non-identity transform.

Example:

import torch
from torch_openreml.covariance import DiagonalMatrix, Adapter
from torch_openreml.covariance.transform import TransformIdentity

# Wrap a 2x2 diagonal matrix so that both variances share a
# single parameter (sum-to-one constraint).
adaptee = DiagonalMatrix(2)

def param_map(params):
    # params[0] drives one variance; the other is 1 - params[0]
    p = torch.sigmoid(params[0])
    return torch.stack([p, 1 - p])

param_specs = {
    "logit": {
        "fixed": False,
        "default": torch.tensor([0.0]),
        "trans": TransformIdentity(),
    }
}

adapter = Adapter(adaptee, param_specs, param_map)
adapter(torch.tensor([0.0]))

tensor([[2.7183, 0.0000],
        [0.0000, 2.7183]])

Methods

`__call__`([free_params])	Build the covariance matrix from the adapter's parameters.
`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 intermediate caching on this adapter and its adaptee.
`enable_cache`()	Enable intermediate caching on this adapter and its adaptee.
`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__()` via the chain rule.
`reset_intermediates`()	Clear the intermediate caches of this adapter and its adaptee.
`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

`adaptee`	The wrapped matrix that this adapter reparameterises.
`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_map`	The mapping function `f(params) -> adaptee_params` that translates the adapter's free parameter tensor to the adaptee's parameter tensor.
`param_names`	Parameter names.
`param_specs`	Parameter specifications, returned as stored rather than copied.
`param_trans`	Parameter transforms.
`repr_dict`	A dictionary representation for display, containing `param_specs`, `adaptee`, and a placeholder for `param_map`.
`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]¶

Build the covariance matrix from the adapter’s parameters.

free_params is checked by build_params(), called with include_fixed=False and trans=False: only the free parameters are kept, and no transform is applied. Those free parameters are passed to param_map, and its output builds the adaptee’s matrix.

Parameters:

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

Returns:

The covariance matrix built by the adaptee.

Return type:

torch.Tensor

Raises:

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

reset_intermediates()[source]¶

Clear the intermediate caches of this adapter and its adaptee.

Extends reset_intermediates(), which clears this adapter’s own cache, by clearing the cache of the wrapped adaptee, whose intermediate results the adapter’s calls produce.

enable_cache()[source]¶

Enable intermediate caching on this adapter and its adaptee.

Extends enable_cache(), which enables caching on this adapter, by enabling it on the wrapped adaptee. A nested adaptee continues the propagation through this same override.

disable_cache()[source]¶

Disable intermediate caching on this adapter and its adaptee.

Extends disable_cache(), which disables caching on this adapter and clears its cache, by disabling caching on the wrapped adaptee and clearing its cache. A nested adaptee continues the propagation through this same override.

manual_grad(free_params=None)[source]¶

Compute the Jacobian of __call__() via the chain rule.

The gradient is obtained by pulling back the adaptee’s gradient through the Jacobian of param_map:

\[\frac{\partial \symbf{V}}{\partial \boldsymbol{\theta}} = \sum_k \frac{\partial \symbf{V}_{\text{adaptee}}}{\partial \phi_k} \cdot \frac{\partial \phi_k}{\partial \boldsymbol{\theta}}\]

where \(\boldsymbol{\phi} = f(\boldsymbol{\theta})\) are the adaptee’s parameters. free_params is checked by build_params(), called with include_fixed=False and trans=False, and the free parameters are then passed to param_map untransformed.

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, after free_params is validated via build_params().
Return type:: tuple

property param_map¶

The mapping function f(params) -> adaptee_params that translates the adapter’s free parameter tensor to the adaptee’s parameter tensor. The parameters are untransformed.

Type:: callable

property adaptee¶

The wrapped matrix that this adapter reparameterises.

Type:: Matrix

property repr_dict¶

A dictionary representation for display, containing param_specs, adaptee, and a placeholder for param_map.

Type:: dict