bartorch.nlop.Parameters

class bartorch.nlop.Parameters(module)

A module’s parameters as one argument of an operator.

FromTorch differentiates by its arguments, so a denoiser whose weights are to be trained inside a graph BART applies has to take them rather than close over them. A module keeps its parameters as several real tensors of several shapes; this is the one complex vector BART can carry, and the way back.

The weights ride in the real part. BART’s operators are complex throughout, so the vector is complex64 and its imaginary half is an exact null direction: nothing reads it, and the gradient that comes back is complex with the answer in its real part – which is what an optimizer over the packed vector wants, and why the packed vector is the thing to hold as a torch.nn.Parameter.

Parameters:

module (torch.nn.Module) – Read for the names and shapes of its parameters. It is not kept differentiably: what the operator differentiates is the vector.

shape

What to give FromTorch as the weights’ shape.

Type:

tuple of int

Examples

>>> weights = nlop.Parameters(denoiser)
>>> prior = nlop.FromTorch(
...     lambda x, w: torch.func.functional_call(denoiser, weights.unpack(w), (x,)),
...     [state, weights.shape],
...     state,
... )
>>> trained = torch.nn.Parameter(weights.pack())
>>> torch.optim.Adam([trained], lr=1e-3)

and afterwards weights.load(trained) puts them back in the module.

__init__(module)

Methods

__init__(module)
load(weights)	Write weights back into the module, after training.
pack()	The module's parameters as they stand, as one complex vector.
unpack(weights)	weights as the mapping torch.func.functional_call takes.