add docstrings

Author: GiovanniCanali

pina/_src/model/block/kan_block.py

@@ -7,7 +7,20 @@
 
 class KANBlock(torch.nn.Module):
     """
-    TODO: docstring.
+    The inner block of the Kolmogorov-Arnold Network (KAN).
+
+    The block applies a spline transformation to the input, optionally combined
+    with a linear transformation of a base activation function. The output is
+    aggregated across input dimensions to produce the final output.
+
+    .. seealso::
+
+        **Original reference**:
+        Liu Z., Wang Y., Vaidya S., Ruehle F., Halverson J., Soljacic M.,
+        Hou T., Tegmark M. (2025).
+        *KAN: Kolmogorov-Arnold Networks*.
+        DOI: `arXiv preprint arXiv:2404.19756.
+        <https://arxiv.org/abs/2404.19756>`_
     """
 
     def __init__(
@@ -119,16 +132,15 @@ def __init__(
 
     def forward(self, x):
         """
-        Forward pass of the :class:`KANBlock`. It transforms the input using a
-        vectorized spline basis and optionally adds a linear transformation of a
-        base activation function.
-
-        The input is expected to have shape (batch_size, input_dimensions) and
-        the output will have shape (batch_size, output_dimensions).
+        Forward pass of the Kolmogorov-Arnold block. The input is passed through
+        the spline transformation, optionally combined with a linear
+        transformation of the base function output, and then aggregated across
+        input dimensions to produce the final output.
 
-        :param torch.Tensor x: The input tensor for the model.
+        :param x: The input tensor for the model.
+        :type x: torch.Tensor | LabelTensor
         :return: The output tensor of the model.
-        :rtype: torch.Tensor
+        :rtype: torch.Tensor | LabelTensor
         """
         y = self.spline(x)
 

pina/_src/model/kolmogorov_arnold_network.py

@@ -5,7 +5,20 @@
 
 class KolmogorovArnoldNetwork(torch.nn.Module):
     """
-    TODO: add docstring.
+    Implementation of Kolmogorov-Arnold Network (KAN).
+
+    The model consists of a sequence of KAN blocks, where each block applies a
+    spline transformation to the input, optionally combined with a linear
+    transformation of a base activation function.
+
+    .. seealso::
+
+        **Original reference**:
+        Liu Z., Wang Y., Vaidya S., Ruehle F., Halverson J., Soljacic M.,
+        Hou T., Tegmark M. (2025).
+        *KAN: Kolmogorov-Arnold Networks*.
+        DOI: `arXiv preprint arXiv:2404.19756.
+        <https://arxiv.org/abs/2404.19756>`_
     """
 
     def __init__(
@@ -78,7 +91,13 @@ def __init__(
 
     def forward(self, x):
         """
-        TODO: add docstring.
+        Forward pass of the KolmogorovArnoldNetwork model. It passes the input
+        through each KAN block in the network and returns the final output.
+
+        :param x: The input tensor for the model.
+        :type x: torch.Tensor | LabelTensor
+        :return: The output tensor of the model.
+        :rtype: torch.Tensor | LabelTensor
         """
         for layer in self.kan_layers:
             x = layer(x)