4.1.5.4 pde.backends.torch.utils module

Defines utilities for the torch backend.

class TorchOperatorType(*args, **kwargs)

Bases: Protocol

An operator that acts on an array.

get_torch_dtype(dtype)

Convert dtype to torch dtype.

Parameters:

dtype (type[Any] | dtype[Any] | _SupportsDType[dtype[Any]] | tuple[Any, Any] | list[Any] | _DTypeDict | str | None | dtype) – dtype which could be a python type, a numpy dtype, or already a torch dtype

Returns:

A proper dtype for torch

Return type:

torch.dtype