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.
`get_intermediates`(params)	Retrieve cached intermediate computation results if still valid.
`grad`([free_params])	Compute the Jacobian of `__call__()` with respect to trainable parameters.
`manual_grad`([free_params])	Compute the Jacobian of `__call__()` with respect to trainable parameters using a closed-form analytic expression.
`map_theta_to_dv`(theta)	An interface compatible with `torch_openreml.REML` that maps parameters to the matrix Jacobian.
`map_theta_to_v`(theta)	An interface compatible with `torch_openreml.REML` that maps parameters to a matrix.
`operands_grad`([free_params])	Compute the Jacobian of each operand with respect to its parameters.
`reset_intermediates`()	Clear the intermediate computation cache.
`set_intermediates`(params, intermediates)	Cache intermediate computation results keyed by parameter hash.
`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.
`operands`	Mapping from operand names to operand matrices or tensors.
`param_defaults`	Parameter defaults.
`param_names`	Parameter names.
`param_specs`	Parameter specifications.
`param_trans`	Parameter transforms.
`repr_dict`	Key-value pairs used to build the string representation.
`shape`	Output matrix shape.

__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 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, 14.7781]],

        [[ 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, 14.7781]]])

grad_names

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