tlm_adjoint.torch

Interface with PyTorch.

Can be used to embed models, differentiated by tlm_adjoint, within a PyTorch calculation. Follows the same principles as described in

• Nacime Bouziani and David A. Ham, ‘Physics-driven machine learning models coupling PyTorch and Firedrake’, 2023, arXiv:2303.06871v3

Module Contents

tlm_adjoint.torch.to_torch_tensors(X, *args, **kwargs)

Convert one or more variables to torch.Tensor objects.

Parameters:

Xvariable or Sequence[variable, …]

Variables to be converted.

args, kwargs

Passed to torch.tensor().

Returns:

tuple[variable, …]

The converted variables.

tlm_adjoint.torch.from_torch_tensors(X, X_t)

Copy data from PyTorch tensors into variables.

Parameters:

Xvariable or Sequence[variable, …]

Output.

X_tSequence[torch.Tensor, …]

Input.

tlm_adjoint.torch.torch_wrapped(forward, space, *, manager=None, clear_caches=True)

Wrap a model, differentiated using tlm_adjoint, so that it can be used with PyTorch.

Parameters:

forwardcallable

Accepts one or more variable arguments, and returns a variable or Sequence of variables.

spacespace or Sequence[space, …]

Defines the spaces for input arguments.

managerEquationManager

Used to create an internal manager via EquationManager.new(). manager() is used if not supplied.

clear_cachesWhether to clear caches before a call of forward.

Returns:

callable

A version of forward with torch.Tensor inputs and outputs.