Neural Network#

class equisolve.nn.Linear(in_tensor: TensorMap, out_tensor: TensorMap, bias: bool = True)[source]#
Parameters:

  • in_tensor – A tensor map that will be accepted in the forward() function. It is used to determine the keys input shape, device and dtype of the input to create linear modules for tensor maps.

  • out_tensor – A tensor map that is used to determine the properties labels and shape of the output tensor map. Because a linear module can change the number of properties, the labels of the properties cannot be persevered.

  • bias – See torch.nn.Linear

forward(tensor: TensorMap) → TensorMap[source]#

Takes a tensor map and applies the modules on each key it.

Parameters:

tensor – input tensor map

classmethod from_module(in_keys: Labels, in_features: int, out_features: int, bias: bool = True, device: device | None = None, dtype: dtype | None = None, many_to_one: bool = True, out_tensor: TensorMap | None = None)[source]#
Parameters:

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

  • in_features – See torch.nn.Linear

  • out_features – See torch.nn.Linear

  • bias – See torch.nn.Linear

  • device – See torch.nn.Linear

  • dtype – See torch.nn.Linear

  • many_to_one – Specifies if a separate module for each block is used. If True the module is deepcopied for each key in the in_keys.

  • out_tensor – A tensor map that is used to determine the properties labels of the output. Because an arbitrary module can change the number of properties, the labels of the properties cannot be persevered. By default the output properties are relabeled using Labels.range.

class equisolve.nn.ModuleTensorMap(module_map: ModuleDict, out_tensor: TensorMap | None = None)[source]#

A wrapper around a torch.nn.ModuleDict to apply each module to the corresponding tensor block in the map using the dict key.

Parameters:

  • module_map – A dictionary of modules with tensor map keys as dict keys each module is applied on a block

  • out_tensor – A tensor map that is used to determine the properties labels of the output. Because an arbitrary module can change the number of properties, the labels of the properties cannot be persevered. By default the output properties are relabeled using Labels.range.

forward(tensor: TensorMap) → TensorMap[source]#

Takes a tensor map and applies the modules on each key it.

Parameters:

tensor – input tensor map

classmethod from_module(in_keys: Labels, module: Module, many_to_one: bool = True, out_tensor: TensorMap | None = None)[source]#

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

Parameters:

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

  • module – The module that is applied on each block.

  • many_to_one – Specifies if a separate module for each block is used. If True the module is deep copied for each key in the in_keys.

  • out_tensor – A tensor map that is used to determine the properties labels of the output. Because an arbitrary module can change the number of properties, the labels of the properties cannot be persevered. By default the output properties are relabeled using Labels.range.

property module_map#

The torch.nn.ModuleDict that maps hashed module keys to a module (see ModuleTensorMap.module_key())

property out_tensor: TensorMap | None#

The tensor map that is used to determine properties labels of the output of forward function.