4.1.6.4 pde.backends.torch.utils module

Defines utilities for the torch backend.

class TorchOperatorBase(*, dtype)

Bases: Module

Base class for operators implemented in torch.

Initialize the torch operator.

Parameters:

dtype (type[Any] | dtype[Any] | _HasDType[dtype[Any]] | _HasNumPyDType[dtype[Any]] | tuple[Any, Any] | list[Any] | _DTypeDict | str) – The data type of the field using the numpy convention

register_array(name, arr)

Register an array as a buffer in the torch module.

Parameters:

• name (str) – The name under which the buffer is registered

• arr (numpy.ndarray or torch.Tensor) – The array to register. If a numpy array is provided, it will be converted to a torch tensor with the appropriate dtype.

Return type:

None

torch_heaviside(x1, x2=None)

Return the Heaviside step function using torch.

This wraps torch.heaviside() and ensures that scalar fallback values are converted to tensors with a dtype compatible with x1.

Parameters:

• x1 (torch.Tensor) – Input values at which the Heaviside function is evaluated.

• x2 (torch.Tensor, optional) – Value used where x1 == 0. If omitted, 0.5 is used.

Returns:

Tensor containing the Heaviside values of x1.

Return type:

torch.Tensor

torch_hypot(x1, x2)

Return the Euclidean norm sqrt(x1**2 + x2**2) using torch.

This wraps torch.hypot() and ensures that both inputs are converted to tensors before evaluation.

Parameters:

• x1 (torch.Tensor) – First input values.

• x2 (torch.Tensor) – Second input values.

Returns:

Tensor containing the element-wise hypotenuse of x1 and x2.

Return type:

torch.Tensor