torch_openreml.covariance.KroneckerProduct¶

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

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.
`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} = \symbf{A} \otimes \symbf{B}\), so its inverse is \(\symbf{V}^{-1} = \symbf{A}^{-1} \otimes \symbf{B}^{-1}\): each operand is inverted through its own inverse(), so the composite is never inverted densely and its \((mp) \times (mp)\) inverse is never formed in full.

Warning

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

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

tensor([[ 0.0233,  0.0000, -0.0108, -0.0000],
        [ 0.0000,  0.0233, -0.0000, -0.0108],
        [-0.0108, -0.0000,  0.0233,  0.0000],
        [-0.0000, -0.0108,  0.0000,  0.0233]])

manual_grad(free_params=None)[source]¶

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