torch_openreml.covariance.Operator

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

Bases: Matrix

Abstract base class for composite covariance matrix operators.

An operator combines one or more Matrix instances and optional fixed torch.Tensor operands into a single covariance matrix. Parameters and gradients are namespaced by operand — a parameter "sigma^2" belonging to operand "residual" is exposed as "residual/sigma^2" in the combined param_names.

Operands may be passed as positional arguments, as keyword arguments, or as a single dictionary. Mixing positional and keyword arguments is not permitted. Positional operands are assigned names "op_0", "op_1", etc.

At least one operand must be a Matrix instance. Pure-tensor operands are treated as fixed matrices with no free parameters.

Every operand is built on a single dtype and device, resolved once for the whole composite:

  1. the input parameters (free_params);

  2. the free-parameter defaults of the Matrix operands, or all parameter defaults when no operand has a free parameter;

  3. the default dtype and device of Torch.

In case 2 the dtype and the device must be equal across all Matrix operands, including any nested inside nested operators, or a ValueError is raised. Matrix operands follow the resolved dtype and device through the parameter tensor they receive, and fixed tensor operands are cast to it.

Initialize the operator from positional or keyword operands.

Parameters:

  • *args – Operands as positional arguments, each a Matrix or torch.Tensor. A single dict argument is also accepted and treated as a named operand mapping.

  • **kwargs – Operands as keyword arguments, mapping operand names to Matrix or torch.Tensor instances. Operand names must be non-empty and must not contain "/".

Raises:

  • ValueError – If both positional and keyword arguments are provided, or if any operand name is empty or contains "/".

  • TypeError – If operands is not a dict, if any operand name is not a string, if any operand is not a Matrix or torch.Tensor, or if no operand is a Matrix.

Example

import torch
from torch_openreml.covariance import Sum, ScalarMatrix

x = Sum(ScalarMatrix(2), ScalarMatrix(2))
x

Sum(
  operands={
    'op_0': ScalarMatrix(shape=(2, 2), param_specs={'sigma^2': {'fixed': False, 'default': tensor([0.]), 'trans': TransformExpPow2()}}),
    'op_1': ScalarMatrix(shape=(2, 2), param_specs={'sigma^2': {'fixed': False, 'default': tensor([0.]), 'trans': TransformExpPow2()}})
  }
)

x = Sum(A = ScalarMatrix(2), B = ScalarMatrix(2))
x

Sum(
  operands={
    'A': ScalarMatrix(shape=(2, 2), param_specs={'sigma^2': {'fixed': False, 'default': tensor([0.]), 'trans': TransformExpPow2()}}),
    'B': ScalarMatrix(shape=(2, 2), param_specs={'sigma^2': {'fixed': False, 'default': tensor([0.]), 'trans': TransformExpPow2()}})
  }
)

x.param_names

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

x(torch.zeros(2))

tensor([[2., 0.],
        [0., 2.]])

x({"A/sigma^2": torch.zeros(1), "B/sigma^2": torch.zeros(1)})

tensor([[2., 0.],
        [0., 2.]])

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

Construct the matrix and return its inverse.

manual_grad([free_params])

Compute the Jacobian of __call__() with respect to free 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().
build_params(free_params=None, include_fixed=True, trans=True, out_format='tensor')[source]

Construct the full parameter tensor by delegating to each operand.

Splits free_params into per-operand slices according to each operand’s num_free_params, calls build_params() on each Matrix operand, and concatenates the results. Fixed tensor operands are skipped.

Parameters:

  • free_params (torch.Tensor or dict) – Flat 1D tensor of free parameters or a dictionary mapping parameter names to tensors. If omitted, default values are used. Default: None.

  • include_fixed (bool, optional) – Whether to include fixed parameters in the output. Passed through to each operand’s build_params(). Default: True.

  • trans (bool, optional) – Whether to apply parameter transforms to the output. Passed through to each operand’s build_params(). Default: True.

  • out_format (str, optional) – Output format. One of "tensor" or "dict". Default: "tensor".

Returns:

Concatenated parameter tensor of length num_params (or num_free_params when include_fixed=False), or a dictionary mapping namespaced parameter names to value tensors.

Return type:

torch.Tensor or dict

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.

  • ValueError – If out_format is not "tensor" or "dict".

import torch
from torch_openreml.covariance import Sum, ScalarMatrix

x = Sum(ScalarMatrix(2), ScalarMatrix(2))
free_params = torch.tensor([0.0, 0.5])

x.build_params(free_params)

tensor([1.0000, 2.7183])

x.build_params()

tensor([1., 1.])

x.build_params(free_params, trans=False)

tensor([0.0000, 0.5000])

x.build_params(free_params, out_format="dict")

{'op_0/sigma^2': tensor([1.]), 'op_1/sigma^2': tensor([2.7183])}
build_operands(free_params=None)[source]

Evaluate each operand at the current free parameters.

Splits free_params into per-operand slices and calls each Matrix operand to produce its matrix. Fixed tensor operands are cast to the dtype and the device of free_params.

Parameters:

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

Returns:

Evaluated operand matrices in the same order as operands.

Return type:

list of torch.Tensor

Example

import torch
from torch_openreml.covariance import Sum, ScalarMatrix

x = Sum(ScalarMatrix(2), ScalarMatrix(2))
v_groups = x.build_operands(torch.tensor([1.0, 2.0]))
print(v_groups[0])
print(v_groups[1])

tensor([[7.3891, 0.0000],
        [0.0000, 7.3891]])
tensor([[54.5982,  0.0000],
        [ 0.0000, 54.5982]])
tree(free_params=None)[source]

Evaluate this operator and every nested operand, returning each result keyed by its path.

Descends the operand structure, calling tree() on nested Operator operands. Paths join the operand names leading to a node with "/", and "/" is the root. Both dictionaries returned share one key set:

  • results maps each path to the matrix the node evaluated to;

  • free_params_by_path maps each path to the free parameter slice the node received, or to None for fixed tensor operands.

The root result is the same matrix __call__() returns.

Parameters:

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

Returns:

(results, free_params_by_path), where results maps paths to evaluated matrices and free_params_by_path maps the same paths to the free parameter slices the corresponding operands received.

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.

Note

Only nested Operator operands are descended into. Every other Matrix appears as a leaf, including Adapter with its adaptee.

Note

Cost: a node’s own result comes from its __call__() while its operands are evaluated separately, so every node below the root is entered twice — once while its parent composes its result, once by the descent. The root is entered once. An operator that caches intermediates serves the second entry from its cache and re-evaluates nothing below it; one that does not cache propagates the re-evaluation down. A preceding __call__() therefore lowers the leaf count under caching operators and changes no value.

Example

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

op = Sum(inner=BlockDiagonal(DiagonalMatrix(2), ScalarMatrix(2)),
         extra=ScalarMatrix(4))
results, free_params_by_path = op.tree(torch.tensor([0.0, 0.5, 1.0, 0.5]))

results

{'/': tensor([[ 3.7183,  0.0000,  0.0000,  0.0000],
         [ 0.0000,  5.4366,  0.0000,  0.0000],
         [ 0.0000,  0.0000, 10.1073,  0.0000],
         [ 0.0000,  0.0000,  0.0000, 10.1073]]),
 'inner': tensor([[1.0000, 0.0000, 0.0000, 0.0000],
         [0.0000, 2.7183, 0.0000, 0.0000],
         [0.0000, 0.0000, 7.3891, 0.0000],
         [0.0000, 0.0000, 0.0000, 7.3891]]),
 'inner/op_0': tensor([[1.0000, 0.0000],
         [0.0000, 2.7183]]),
 'inner/op_1': tensor([[7.3891, 0.0000],
         [0.0000, 7.3891]]),
 'extra': tensor([[2.7183, 0.0000, 0.0000, 0.0000],
         [0.0000, 2.7183, 0.0000, 0.0000],
         [0.0000, 0.0000, 2.7183, 0.0000],
         [0.0000, 0.0000, 0.0000, 2.7183]])}

results["inner/op_0"]

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

free_params_by_path["inner"]

tensor([0.0000, 0.5000, 1.0000])
grad_tree(free_params=None)[source]

Compute the gradient of this operator and of every nested operand, returning each keyed by its path.

Descends the operand structure, calling grad_tree() on nested Operator operands. Paths join the operand names leading to a node with "/", and "/" is the root — the same keys tree() produces, in the same order. Both dictionaries returned share one key set:

  • grads maps each path to the (grad, grad_names) pair grad() returns for that node, or to (None, []) when the node has no free parameters;

  • free_params_by_path maps each path to the free parameter slice the node received, or to None for fixed tensor operands.

Each node’s grad has shape (num_free_params, *shape) for that node’s own free parameters and its own matrix shape, so a nested node’s Jacobian is not embedded in its parent’s shape. The root entry is the same pair grad() returns for the whole composite.

grad_names are namespaced from the node that produced them, as grad() returns them: at the root they are the composite’s full names, and below it joining the path to a name with "/" gives the name the root exposes in free_param_names.

Parameters:

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

Returns:

(grads, free_params_by_path), where grads maps paths to (grad, grad_names) pairs and free_params_by_path maps the same paths to the free parameter slices the corresponding operands received.

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.

Note

Only nested Operator operands are descended into. Every other Matrix appears as a leaf, including Adapter with its adaptee.

Note

Cost: every Jacobian is computed independently, so a node at depth d is evaluated d + 1 times — once for the root’s Jacobian and once for each intermediate ancestor’s, plus once for its own entry. No Jacobian is cached, so a preceding grad() changes nothing. Viewing a deep composite in full is correspondingly expensive.

Note

A node’s Jacobian is in that node’s own shape, so the tree does not show how a child’s Jacobian is placed into its parent’s — how BlockDiagonal pads it into a block, for instance. Only the root entry carries the composite’s full shape.

Example

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

op = Sum(inner=BlockDiagonal(DiagonalMatrix(2), ScalarMatrix(2)),
         extra=ScalarMatrix(4))
grads, free_params_by_path = op.grad_tree(torch.tensor([0.0, 0.5, 1.0, 0.5]))

grads["inner/op_0"]

(tensor([[[2.0000, 0.0000],
          [0.0000, 0.0000]],
 
         [[0.0000, 0.0000],
          [0.0000, 5.4366]]]),
 ['sigma^2_0', 'sigma^2_1'])

grads["inner"][1]

['op_0/sigma^2_0', 'op_0/sigma^2_1', 'op_1/sigma^2']
operands_grad(free_params=None)[source]

Compute the Jacobian of each operand with respect to its parameters.

Splits free_params into per-operand slices, calls grad() on each Matrix operand, and prefixes the returned names with the operand name. Fixed tensor operands contribute None and an empty name list.

Parameters:

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

Returns:

(grad_groups, grad_name_groups), where grad_groups is a list of per-operand Jacobian tensors or None for fixed operands, and grad_name_groups is a list of corresponding namespaced parameter name lists.

Return type:

tuple

Example

import torch
from torch_openreml.covariance import Sum, ScalarMatrix

x = Sum(ScalarMatrix(2), ScalarMatrix(2))
grad_groups, grad_name_groups = x.operands_grad(torch.tensor([1.0, 2.0]))
print(grad_groups[0])
print(grad_groups[1])
print(grad_name_groups[0])
print(grad_name_groups[1])

tensor([[[14.7781,  0.0000],
         [ 0.0000, 14.7781]]])
tensor([[[109.1963,   0.0000],
         [  0.0000, 109.1963]]])
['op_0/sigma^2']
['op_1/sigma^2']
operands_inverse(free_params=None)[source]

Compute the inverse of each operand at the current free parameters.

Splits free_params into per-operand slices, calls inverse() on each Matrix operand, and casts and inverts each fixed tensor operand. A Matrix operand is never built here, so an analytic inverse — or a nested operator’s own structured inverse — is used instead of inverting the operand densely.

Parameters:

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

Returns:

Each operand’s inverse, in the order of operands.

Return type:

list of torch.Tensor

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, Sum

x = Sum(ScalarMatrix(2), ScalarMatrix(2))
inverses = x.operands_inverse(torch.tensor([1.0, 2.0]))
print(inverses[0])
print(inverses[1])

tensor([[0.1353, 0.0000],
        [0.0000, 0.1353]])
tensor([[0.0183, 0.0000],
        [0.0000, 0.0183]])
reset_intermediates()[source]

Clear the intermediate caches of this operator and its operands.

Extends reset_intermediates(), which clears this operator’s own cache, by recursively clearing the cache of every Matrix operand. Nested operators continue the descent through this same override, so resetting a node resets everything below it.

enable_cache()[source]

Enable intermediate caching on this operator and its operands.

Extends enable_cache(), which enables caching on this operator, by recursively enabling it on every Matrix operand. Nested operators continue the descent through this same override, so enabling a node enables everything below it.

disable_cache()[source]

Disable intermediate caching on this operator and its operands.

Extends disable_cache(), which disables caching on this operator and clears its cache, by recursively disabling caching on every Matrix operand and clearing their caches. Nested operators continue the descent through this same override, so disabling a node disables everything below it.

property operands

Mapping from operand names to operand matrices or tensors.

Type:

dict

property param_specs

Parameter specifications, returned as stored rather than copied.

A specification edited in place (for example setting "fixed" to True) therefore takes effect on the next call.

Type:

dict

property repr_dict

Key-value pairs used to build the string representation.

Type:

dict