torch_openreml.covariance.KroneckerProduct

class torch_openreml.covariance.KroneckerProduct(*args, **kwargs)

Bases: Operator

Kronecker product of two covariance matrices.

\[\symbf{V} = \symbf{A} \otimes \symbf{B}\]

If \(\symbf{A}\) is \(m \times m\) and \(\symbf{B}\) is \(n \times n\), the result is an \(mn \times mn\) matrix. Either or both operands may be trainable Matrix instances or fixed torch.Tensor values.

Initialize a Kronecker product operator from exactly two operands.

Parameters:

• *args – Exactly two operands as positional arguments or a single dict. The first is \(\symbf{A}\), the second \(\symbf{B}\).

• **kwargs – Exactly two operands as keyword arguments.

Raises:

ValueError – If the number of operands is not exactly two.

Example:

import torch
from torch_openreml.covariance import AR1Matrix, ScalarMatrix, KroneckerProduct

op = KroneckerProduct(time=AR1Matrix(2), subject=ScalarMatrix(2))
params = torch.tensor([1.0, 1.0, 1.0])
op(params)

tensor([[54.5982,  0.0000, 25.2307,  0.0000],
        [ 0.0000, 54.5982,  0.0000, 25.2307],
        [25.2307,  0.0000, 54.5982,  0.0000],
        [ 0.0000, 25.2307,  0.0000, 54.5982]])

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)

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)

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

Applies the product rule: if \(\symbf{V} = \symbf{A} \otimes \symbf{B}\) then the gradient with respect to \(\theta_{\symbf{A}}\) is \(\frac{\partial \symbf{A}}{\partial \theta_{\symbf{A}}} \otimes \symbf{B}\), and similarly for \(\theta_{\symbf{B}}\). Per-operand Jacobians from operands_grad() are Kronecker-multiplied by the other operand’s value.

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 AR1Matrix, ScalarMatrix, KroneckerProduct

op = KroneckerProduct(time=AR1Matrix(2), subject=ScalarMatrix(2))
params = torch.tensor([1.0, 1.0, 1.0])
grad, grad_names = op.manual_grad(params)
grad

tensor([[[109.1963,   0.0000,  50.4615,   0.0000],
         [  0.0000, 109.1963,   0.0000,  50.4615],
         [ 50.4615,   0.0000, 109.1963,   0.0000],
         [  0.0000,  50.4615,   0.0000, 109.1963]],

        [[  0.0000,   0.0000,  21.4693,   0.0000],
         [  0.0000,   0.0000,   0.0000,  21.4693],
         [ 21.4693,   0.0000,   0.0000,   0.0000],
         [  0.0000,  21.4693,   0.0000,   0.0000]],

        [[109.1963,   0.0000,  50.4615,   0.0000],
         [  0.0000, 109.1963,   0.0000,  50.4615],
         [ 50.4615,   0.0000, 109.1963,   0.0000],
         [  0.0000,  50.4615,   0.0000, 109.1963]]])

grad_names

['time/sigma^2', 'time/rho', 'subject/sigma^2']