torch_openreml.covariance.BlockDiagonal

class torch_openreml.covariance.BlockDiagonal(*args, **kwargs)[source]

Bases: Operator

Block diagonal covariance matrix formed from two or more operands.

\[\symbf{V} = \mathrm{blockdiag}(\symbf{V}_0, \symbf{V}_1, \ldots)\]

Each operand contributes a contiguous block along the diagonal. Parameters and gradients are namespaced by operand name and aggregated into a single joint parameter tensor, following the convention of Operator.

Initialize a block diagonal operator from two or more operands.

Parameters:

  • *args – Two or more operands as positional arguments, each a Matrix or torch.Tensor. A single dict argument is also accepted.

  • **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 ScalarMatrix, DiagonalMatrix, BlockDiagonal

block = BlockDiagonal(
    residual=ScalarMatrix(3),
    random=DiagonalMatrix(2)
)
free_params = torch.tensor([0.5, 0.0, 1.0])
block(free_params)

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

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])

Return the inverse of the composite matrix from its closed form.

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)[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

inverse(free_params=None)[source]

Return the inverse of the composite matrix from its closed form.

The composite is \(\symbf{V} = \mathrm{blockdiag}(\symbf{V}_0, \symbf{V}_1, \ldots)\), so its inverse is \(\symbf{V}^{-1} = \mathrm{blockdiag}(\symbf{V}_0^{-1}, \symbf{V}_1^{-1}, \ldots)\): each operand is inverted through its own inverse() and the results are placed back on the diagonal, so the composite is never inverted densely.

Warning

Every operand must be square, otherwise the composite has no inverse and raises.

Parameters:

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

Returns:

The n x n inverse of the composite matrix.

Return type:

torch.Tensor

Example:

import torch
from torch_openreml.covariance import ScalarMatrix, DiagonalMatrix, BlockDiagonal

op = BlockDiagonal(
    residual=ScalarMatrix(3),
    random=DiagonalMatrix(2)
)
free_params = torch.tensor([0.5, 0.0, 1.0])
op.inverse(free_params)

tensor([[0.3679, 0.0000, 0.0000, 0.0000, 0.0000],
        [0.0000, 0.3679, 0.0000, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.3679, 0.0000, 0.0000],
        [0.0000, 0.0000, 0.0000, 1.0000, 0.0000],
        [0.0000, 0.0000, 0.0000, 0.0000, 0.1353]])
manual_grad(free_params=None)[source]

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

Evaluates each operand’s gradient via operands_grad(), then places each per-operand Jacobian into the corresponding block of a zero-initialised full Jacobian matching the block diagonal output shape. Non-matrix operands contribute zero blocks.

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 ScalarMatrix, DiagonalMatrix, BlockDiagonal

block = BlockDiagonal(
    A=ScalarMatrix(3),
    B=DiagonalMatrix(2)
)
free_params = torch.tensor([0.5, 0.0, 1.0])
grad, grad_names = block.manual_grad(free_params)
grad

tensor([[[ 5.4366,  0.0000,  0.0000,  0.0000,  0.0000],
         [ 0.0000,  5.4366,  0.0000,  0.0000,  0.0000],
         [ 0.0000,  0.0000,  5.4366,  0.0000,  0.0000],
         [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
         [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000]],

        [[ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
         [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
         [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000],
         [ 0.0000,  0.0000,  0.0000,  2.0000,  0.0000],
         [ 0.0000,  0.0000,  0.0000,  0.0000,  0.0000]],

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

grad_names

['A/sigma^2', 'B/sigma^2_0', 'B/sigma^2_1']