fix doc model part 2

pina/model/block/average_neural_operator_block.py

@@ -1,4 +1,4 @@
-"""Module for Averaging Neural Operator Layer class."""
+"""Module for the Averaging Neural Operator Block class."""
 
 import torch
 from torch import nn
@@ -7,12 +7,12 @@ from ...utils import check_consistency
 
 class AVNOBlock(nn.Module):
     r"""
-    The PINA implementation of the inner layer of the Averaging Neural Operator.
+    The inner block of the Averaging Neural Operator.
 
     The operator layer performs an affine transformation where the convolution
     is approximated with a local average. Given the input function
-    :math:`v(x)\in\mathbb{R}^{\rm{emb}}` the layer computes
-    the operator update :math:`K(v)` as:
+    :math:`v(x)\in\mathbb{R}^{\rm{emb}}` the layer computes the operator update
+    :math:`K(v)` as:
 
     .. math::
         K(v) = \sigma\left(Wv(x) + b + \frac{1}{|\mathcal{A}|}\int v(y)dy\right)
@@ -28,18 +28,20 @@ class AVNOBlock(nn.Module):
 
     .. seealso::
 
-        **Original reference**: Lanthaler S. Li, Z., Kovachki,
-        Stuart, A. (2020). *The Nonlocal Neural Operator: Universal
-        Approximation*.
+        **Original reference**: Lanthaler S., Li, Z., Stuart, A. (2020).
+        *The Nonlocal Neural Operator: Universal Approximation*.
         DOI: `arXiv preprint arXiv:2304.13221.
         <https://arxiv.org/abs/2304.13221>`_
-
     """
 
     def __init__(self, hidden_size=100, func=nn.GELU):
         """
-        :param int hidden_size: Size of the hidden layer, defaults to 100.
-        :param func: The activation function, default to nn.GELU.
+        Initialization of the :class:`AVNOBlock` class.
+
+        :param int hidden_size: The size of the hidden layer.
+            Defaults is ``100``.
+        :param func: The activation function.
+            Default is :class:`torch.nn.GELU`.
         """
         super().__init__()
 
@@ -52,17 +54,11 @@ class AVNOBlock(nn.Module):
 
     def forward(self, x):
         r"""
-        Forward pass of the layer, it performs a sum of local average
-        and an affine transformation of the field.
+        Forward pass of the block. It performs a sum of local average and an
+        affine transformation of the field.
 
-        :param torch.Tensor x: The input tensor for performing the
-            computation. It expects a tensor :math:`B \times N \times D`,
-            where :math:`B` is the batch_size, :math:`N` the number of points
-            in the mesh, :math:`D` the dimension of the problem. In particular
-            :math:`D` is the codomain of the function :math:`v`. For example
-            a scalar function has :math:`D=1`, a 4-dimensional vector function
-            :math:`D=4`.
-        :return: The output tensor obtained from Average Neural Operator Block.
+        :param torch.Tensor x: The input tensor for performing the computation.
+        :return: The output tensor.
         :rtype: torch.Tensor
         """
         return self._func(self._nn(x) + torch.mean(x, dim=1, keepdim=True))