bartorch.nlop.FromTorch

bartorch.nlop.FromTorch#

class bartorch.nlop.FromTorch(fn, ishape, oshape)#

A nonlinear operator from a differentiable torch function.

The derivative is torch’s forward-mode Jacobian-vector product, and its adjoint the reverse-mode vector-Jacobian product, both at the last evaluated point.

Parameters:

  • fn (callable) – A differentiable function on tensors. With several input shapes it takes that many arguments, and it may return several tensors.

  • ishape (tuple of int, or a sequence of them) – Domain and codomain, C order.

  • oshape (tuple of int, or a sequence of them) – Domain and codomain, C order.

Examples

>>> F = FromTorch(lambda p: p[0] * torch.exp(-t / p[1]), (2,), t.shape)
>>> optim.IRGNM()(measured, F, x0=torch.tensor([1.0, 20.0]))

A denoiser whose weights are an argument, so that they train through a graph BART applies:

>>> prior = FromTorch(lambda x, w: w * x, [state, ()], state)
>>> nlop.chain(cell, prior, output=0, input=0)

Notes

A real argument is carried in the real part of a complex one: BART’s operators are complex throughout, so a weight arrives as w + 0j and its gradient comes back complex. Take the real part of it.

__init__(fn, ishape, oshape)#

Methods

__init__(fn, ishape, oshape)

adjoint(dy)

DF(x)^H dy at the last evaluated point, for one input and one output.

chain(other, *[, output, input])

Output output of self into input input of other.

combine(other)

self and other side by side, sharing nothing.

del_out([output])

Drop an output, and everything computed only for it.

derivative(dx)

DF(x) dx at the last evaluated point, for one input and one output.

dup([a, b])

Make two inputs of the same shape one input, kept at a.

flatten([inputs_only])

Every input as one flat vector, and every output as another.

forward(*xs)

F(x), which also fixes where every derivative is taken.

jacobian([output, input])

DF/dx_input of one output, as a LinearOperator.

linearize(*xs)

The derivative at x, as a LinearOperator.

link([output, input])

Tie an output back into an input; both arguments go away.

permute_inputs(perm)

Reorder the inputs: the new input i is the old perm[i].

permute_outputs(perm)

Reorder the outputs: the new output o is the old perm[o].

pin(input, value)

Fix one input to value; the input goes away.

reshape_input(input, shape)

This operator with one input's shape written differently.

reshape_output(output, shape)

This operator with one output's shape written differently.

shift_input(new, old)

Move one input to another position, the rest closing up behind it.

shift_output(new, old)

Move one output to another position, the rest closing up behind it.

stack_inputs(a, b, axis)

Make two inputs one, concatenated along a C-order axis.

stack_outputs(a, b, axis)

Make two outputs one, concatenated along a C-order axis.

Attributes

device

ishape

The domain, for an operator with one input.

ishapes

The shape of each argument, outputs and inputs kept apart.

oshape

The codomain, for an operator with one output.

oshapes