Neural Network#

Modules#

class metatensor.learn.nn.ModuleMap(in_keys: Labels, modules: List[Module], out_properties: List[Labels] | None = None)[source]#

A class that imitates torch.nn.ModuleDict. In its forward function the module at position i given on construction by :param modules: is applied to the tensor block that corresponding to the`i`th key in :param in_keys:.

Parameters:

  • in_keys (Labels) – A metatensor.Labels object with the keys of the module map that are assumed to be in the input tensor map in the forward() function.

  • modules (List[Module]) – A sequence of modules applied in the forward() function on the input TensorMap. Each module corresponds to one LabelsEntry in :param in_keys: that determines on which TensorBlock the module is applied on. :param modules: and :param in_keys: must match in length.

  • out_properties (List[Labels] | None) –

    A list of labels that is used to determine the properties labels of the output. Because a module could change the number of properties, the labels of the properties cannot be persevered. By default the output properties are relabeled using Labels.range with “_” as key.

    >>> import torch
>>> import numpy as np
>>> from copy import deepcopy
>>> from metatensor import Labels, TensorBlock, TensorMap
>>> from metatensor.learn.nn import ModuleMap

    Create simple block

    >>> block_1 = TensorBlock(
...     values=torch.tensor(
...         [
...             [1.0, 2.0, 4.0],
...             [3.0, 5.0, 6.0],
...         ]
...     ),
...     samples=Labels(
...         ["system", "atom"],
...         np.array(
...             [
...                 [0, 0],
...                 [0, 1],
...             ]
...         ),
...     ),
...     components=[],
...     properties=Labels(["properties"], np.array([[0], [1], [2]])),
... )
>>> block_2 = TensorBlock(
...     values=torch.tensor(
...         [
...             [5.0, 8.0, 2.0],
...             [1.0, 2.0, 8.0],
...         ]
...     ),
...     samples=Labels(
...         ["system", "atom"],
...         np.array(
...             [
...                 [0, 0],
...                 [0, 1],
...             ]
...         ),
...     ),
...     components=[],
...     properties=Labels(["properties"], np.array([[3], [4], [5]])),
... )
>>> keys = Labels(names=["key"], values=np.array([[0], [1]]))
>>> tensor = TensorMap(keys, [block_1, block_2])

    Create modules

    >>> linear = torch.nn.Linear(3, 1, bias=False)
>>> with torch.no_grad():
...     _ = linear.weight.copy_(torch.tensor([1.0, 1.0, 1.0]))
...
>>> modules = [linear, deepcopy(linear)]
>>> # you could also extend the module by some nonlinear activation function

    Create ModuleMap from this ModucDict and apply it

    >>> module_map = ModuleMap(tensor.keys, modules)
>>> out = module_map(tensor)
>>> out
TensorMap with 2 blocks
keys: key
       0
       1
>>> out[0].values
tensor([[ 7.],
        [14.]], grad_fn=<MmBackward0>)
>>> out[1].values
tensor([[15.],
        [11.]], grad_fn=<MmBackward0>)

    Let’s look at the metadata

    >>> tensor[0]
TensorBlock
    samples (2): ['system', 'atom']
    components (): []
    properties (3): ['properties']
    gradients: None
>>> out[0]
TensorBlock
    samples (2): ['system', 'atom']
    components (): []
    properties (1): ['_']
    gradients: None

    It got completely lost because we cannot know in general what the output is. You can add in the initialization of the ModuleMap a TensorMap that contains the intended output Labels.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

classmethod from_module(in_keys: Labels, module: Module, out_properties: List[Labels] | None = None)[source]#

A wrapper around one torch.nn.Module applying the same type of module on each tensor block.

Parameters:

  • in_keys (Labels) – A metatensor.Labels object that determines the keys of the module map that are ass the TensorMaps that are assumed to be in the input tensor map in the forward() function.

  • module (Module) – The module that is applied on each block.

  • out_properties (List[Labels] | None) – A list of labels that is used to determine the properties labels of the output. Because a module could change the number of properties, the labels of the properties cannot be persevered. By default the output properties are relabeled using Labels.range with “_” as key.

>>> import torch
>>> import numpy as np
>>> from metatensor import Labels, TensorBlock, TensorMap
>>> block_1 = TensorBlock(
...     values=torch.tensor(
...         [
...             [1.0, 2.0, 4.0],
...             [3.0, 5.0, 6.0],
...         ]
...     ),
...     samples=Labels(
...         ["system", "atom"],
...         np.array(
...             [
...                 [0, 0],
...                 [0, 1],
...             ]
...         ),
...     ),
...     components=[],
...     properties=Labels(["properties"], np.array([[0], [1], [2]])),
... )
>>> block_2 = TensorBlock(
...     values=torch.tensor(
...         [
...             [5.0, 8.0, 2.0],
...             [1.0, 2.0, 8.0],
...         ]
...     ),
...     samples=Labels(
...         ["system", "atom"],
...         np.array(
...             [
...                 [0, 0],
...                 [0, 1],
...             ]
...         ),
...     ),
...     components=[],
...     properties=Labels(["properties"], np.array([[0], [1], [2]])),
... )
>>> keys = Labels(names=["key"], values=np.array([[0], [1]]))
>>> tensor = TensorMap(keys, [block_1, block_2])
>>> linear = torch.nn.Linear(3, 1, bias=False)
>>> with torch.no_grad():
...     _ = linear.weight.copy_(torch.tensor([1.0, 1.0, 1.0]))
...
>>> # you could also extend the module by some nonlinear activation function
>>> from metatensor.learn.nn import ModuleMap
>>> module_map = ModuleMap.from_module(tensor.keys, linear)
>>> out = module_map(tensor)
>>> out[0].values
tensor([[ 7.],
        [14.]], grad_fn=<MmBackward0>)
>>> out[1].values
tensor([[15.],
        [11.]], grad_fn=<MmBackward0>)
forward(tensor: TensorMap) → TensorMap[source]#

Apply the modules on each block in tensor. tensor must have the same set of keys as the modules used to initialize this ModuleMap.

Parameters:

tensor (TensorMap) – input tensor map

Return type:

TensorMap

get_module(key: LabelsEntry)[source]#
Parameters:

key (LabelsEntry) – key of module which should be returned

Return module:

returns he torch.nn.Module corresponding to the :param key:

property in_keys: Labels#

A list of labels that defines the initialized keys with corresponding modules of this module map.

property out_properties: None | List[Labels]#

A list of labels that is used to determine properties labels of the output of forward function.

repr_as_module_dict() → str[source]#

Returns a string that is easier to read that the standard __repr__ showing the mapping from label entry key to module.

Return type:

str

class metatensor.learn.nn.Linear(in_keys: Labels, in_features: int | List[int], out_features: int | List[int], bias: bool | List[bool] = True, device: device | None = None, dtype: dtype | None = None, out_properties: List[Labels] | None = None)[source]#

Construct a module map with only linear modules.

Parameters:

  • in_keys (Labels) – The keys that are assumed to be in the input tensor map in the forward() function.

  • in_features (int | List[int]) – Specifies the dimensionality of the input after it has been applied on the input tensor map. If a list of integers is given, it specifies the input dimension for each block, therefore it should have the same length as :param in_keys:. If only one integer is given, it is assumed that the same be applied on each block.

  • out_features (int | List[int]) – Specifies the dimensionality of the output after it has been applied on the input tensor map. If a list of integers is given, it specifies the output dimension for each block, therefore it should have the same length as :param in_keys:. If only one integer is given, it is assumed that the same be applied on each block.

  • bias (bool | List[bool]) – Specifies if a bias term (offset) should be applied on the input tensor map in the forward function. If a list of bools is given, it specifies the bias term for each block, therefore it should have the same length as :param in_keys:. If only one bool value is given, it is assumed that the same be applied on each block.

  • device (device | None) – Specifies the torch device of the values. If None the default torch device is taken.

  • dtype (dtype | None) – Specifies the torch dtype of the values. If None the default torch dtype is taken.

  • out_properties (List[Labels] | None) – A list of labels that is used to determine the properties labels of the output. Because a module could change the number of properties, the labels of the properties cannot be persevered. By default the output properties are relabeled using Labels.range.

Initialize internal Module state, shared by both nn.Module and ScriptModule.

classmethod from_weights(weights: TensorMap, bias: TensorMap | None = None)[source]#

Construct a linear module map from a tensor map for the weights and bias.

Parameters:

  • weights (TensorMap) – The weight tensor map from which we create the linear modules. The properties of the tensor map describe the input dimension and the samples describe the output dimension.

  • bias (TensorMap | None) – The weight tensor map from which we create the linear layers.