Type annotations: PyTorch

PyTorch-related type annotations for this project.

class tad_multicharge.typing.pytorch.DD

Collection of torch.device and torch.dtype.

device

Device on which a tensor lives.

dtype

Floating point precision of a tensor.

class tad_multicharge.typing.pytorch.ModuleLike

nn.Module with TensorLike-style helpers.

property allowed_dtypes

Collection of dtypes supported by the model.

property dd

Shortcut combining device and dtype.

property device

Device inferred from the first registered buffer.

property dtype

Floating point dtype inferred from the first registered buffer.

to(*args, **kwargs)

Move and/or cast the parameters and buffers.

This can be called as

to(device=None, dtype=None, non_blocking=False)

to(dtype, non_blocking=False)

to(tensor, non_blocking=False)

to(memory_format=torch.channels_last)

Its signature is similar to torch.Tensor.to(), but only accepts floating point or complex dtypes. In addition, this method will only cast the floating point or complex parameters and buffers to dtype (if given). The integral parameters and buffers will be moved device, if that is given, but with dtypes unchanged. When non_blocking is set, it tries to convert/move asynchronously with respect to the host if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices.

See below for examples.

Note

This method modifies the module in-place.

Args:

device (torch.device): the desired device of the parameters

and buffers in this module

dtype (torch.dtype): the desired floating point or complex dtype of

the parameters and buffers in this module

tensor (torch.Tensor): Tensor whose dtype and device are the desired

dtype and device for all parameters and buffers in this module

memory_format (torch.memory_format): the desired memory

format for 4D parameters and buffers in this module (keyword only argument)

Returns:

Module: self

Examples:

>>> # xdoctest: +IGNORE_WANT("non-deterministic")
>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)

type(dst_type)

Casts all parameters and buffers to dst_type.

Note

This method modifies the module in-place.

Args:

dst_type (type or string): the desired type

Returns:

Module: self

class tad_multicharge.typing.pytorch.Molecule

Representation of fundamental molecular structure (atom types and postions).

numbers

Tensor of atomic numbers

positions

Tensor of 3D coordinates of shape (n, 3)

class tad_multicharge.typing.pytorch.TensorLike(device=None, dtype=None)

Provide device and dtype as well as torch.Tensor.to() and torch.Tensor.type() for other classes.

The selection of torch.Tensor variables to change within the class is handled by searching __slots__. Hence, if one wants to use this functionality the subclass of TensorLike must specify __slots__.

property allowed_dtypes

Specification of dtypes that the TensorLike object can take. Defaults to float types and must be overridden by subclass if float are not allowed. The IndexHelper is an example that should only allow integers.

Returns:

Collection of allowed dtypes the TensorLike object can take.

Return type:

tuple[torch.dtype, …]

cpu()

Returns a copy of the TensorLike instance on the CPU.

This method creates and returns a new copy of the TensorLike instance on the CPU.

Returns:

A copy of the TensorLike instance placed on the CPU.

Return type:

TensorLike

property dd

Shortcut for device and dtype.

property device

The device on which the class object resides.

property dtype

Floating point dtype used by class object.

override_device(device)

Override the device of the class object.

Warning

This does not change the device of the underlying tensors. It only changes the device of the class object. Use with caution.

Parameters:

device (torch.device) – Device to override the current device.

override_dtype(dtype)

Override the dtype of the class object.

Warning

This does not change the dtype of the underlying tensors. It only changes the dtype of the class object. Use with caution.

Parameters:

dtype (torch.dtype) – Floating point dtype to override the current dtype.

to(device=None, dtype=None)

Returns a copy of the TensorLike instance on the specified device.

This method creates and returns a new copy of the TensorLike instance on the specified device “device”.

Parameters:

device (torch.device) – Device to which all associated tensors should be moved.

Returns:

A copy of the TensorLike instance placed on the specified device.

Return type:

TensorLike

Notes

If the TensorLike instance is already on the desired device self will be returned.

type(dtype)

Returns a copy of the TensorLike instance with specified floating point type. This method creates and returns a new copy of the TensorLike instance with the specified dtype.

Parameters:

dtype (torch.dtype) – Floating point type.

Returns:

A copy of the TensorLike instance with the specified dtype.

Return type:

TensorLike

Notes

If the TensorLike instance has already the desired dtype Self will be returned.

tad_multicharge.typing.pytorch.get_default_device()

Default device for tensors.

Returns:

PyTorch device type.

Return type:

torch.device

tad_multicharge.typing.pytorch.get_default_dtype()

Default data type for floating point tensors.

Returns:

PyTorch dtype type.

Return type:

torch.dtype