add description parameter for all node objects

src/caskade/backend.py

@@ -198,10 +198,10 @@
         return self.module.all(array)
 
     def log(self, array):
         return self.module.log(array)
 
     def exp(self, array):
         return self.module.exp(array)
 
     def sum(self, array, axis=None):
         return self.module.sum(array, axis=axis)
@@ -213,7 +213,7 @@
         return self.jax.nn.sigmoid(array)
 
     def _sigmoid_numpy(self, array):
         return 1 / (1 + self.module.exp(-array))
 
     def _logit_torch(self, array):
         return self.module.logit(array)
@@ -224,8 +224,5 @@
     def _logit_numpy(self, array):
         return np.log(array / (1 - array))
 
-    def sqrt(self, array):
-        return self.module.sqrt(array)
-
 
 backend = Backend()

src/caskade/base.py

@@ -61,6 +61,7 @@ def __init__(
         self,
         name: Optional[str] = None,
         link: Optional[Union["Node", tuple["Node"]]] = None,
+        description: str = "",
     ):
         if name is None:
             name = self.__class__.__name__
@@ -73,6 +74,7 @@ def __init__(
         self._parents = set()
         self._active = False
         self._type = "node"
+        self.description = description
         self.meta = meta()
         self.saveattrs = set()
         if link is not None:

src/caskade/module.py

@@ -67,8 +67,8 @@ def otherfun(self, x, c = None):
     )  # These tuples will not be converted to NodeTuple objects
     graphviz_types = {"module": {"style": "solid", "color": "black", "shape": "ellipse"}}
 
-    def __init__(self, name: Optional[str] = None):
-        super().__init__(name=name)
+    def __init__(self, name: Optional[str] = None, **kwargs):
+        super().__init__(name=name, **kwargs)
         self.dynamic_params = ()
         self.all_dynamic_value = True
         self.pointer_params = ()

src/caskade/param.py

@@ -103,8 +103,9 @@ def __init__(
         dynamic_value: Optional[Union[ArrayLike, float, int]] = None,
         dtype: Optional[Any] = None,
         device: Optional[Any] = None,
+        **kwargs,
     ):
-        super().__init__(name=name)
+        super().__init__(name=name, **kwargs)
         if value is not None and dynamic_value is not None:
             raise ParamConfigurationError("Cannot set both value and dynamic value")
         if isinstance(value, dynamic):