add support for TensorMap as argument for bias

Author: agoscinski

src/equisolve/nn/module_tensor.py

@@ -7,7 +7,7 @@
     from metatensor import Labels, LabelsEntry, TensorBlock, TensorMap
 
 from copy import deepcopy
-from typing import List, Optional
+from typing import List, Optional, Union
 
 import torch
 from torch.nn import Module, ModuleDict
@@ -167,23 +167,36 @@ class Linear(ModuleTensorMap):
         properties, the labels of the properties cannot be persevered.
 
     :param bias:
-        See :py:class:`torch.nn.Linear`
+        See :py:class:`torch.nn.Linear` for bool as input. For each TensorMap key the
+        bias can be also individually tuend by using a TensorMap with one value for the
+        bool.
     """
 
     def __init__(
         self,
         in_tensor: TensorMap,
         out_tensor: TensorMap,
-        bias: bool = True,
+        bias: Union[bool, TensorMap] = True,
     ):
+        if isinstance(bias, bool):
+            blocks = [
+                TensorBlock(
+                    values=torch.tensor(bias).reshape(1, 1),
+                    samples=Labels.range("_", 1),
+                    components=[],
+                    properties=Labels.range("_", 1),
+                )
+                for _ in in_tensor.keys
+            ]
+            bias = TensorMap(keys=in_tensor.keys, blocks=blocks)
         module_map = ModuleDict()
         for key, in_block in in_tensor.items():
             module_key = ModuleTensorMap.module_key(key)
             out_block = out_tensor.block(key)
             module = torch.nn.Linear(
                 len(in_block.properties),
                 len(out_block.properties),
-                bias,
+                bias.block(key).values.flatten()[0],
                 in_block.values.device,
                 in_block.values.dtype,
             )