From 18682f8155292bd8e0673ea7bd48d0895f956249 Mon Sep 17 00:00:00 2001
From: Laurent <laurent@lagon.tech>
Date: Thu, 1 Feb 2024 21:46:23 +0000
Subject: [PATCH] (doc/fluxion/maxpool) add/convert docstrings to mkdocstrings
 format

---
 src/refiners/fluxion/layers/maxpool.py | 50 +++++++++++++++++++++++---
 1 file changed, 46 insertions(+), 4 deletions(-)

diff --git a/src/refiners/fluxion/layers/maxpool.py b/src/refiners/fluxion/layers/maxpool.py
index 060f210..e62c89a 100644
--- a/src/refiners/fluxion/layers/maxpool.py
+++ b/src/refiners/fluxion/layers/maxpool.py
@@ -1,9 +1,20 @@
-from torch import nn
+from torch.nn import MaxPool1d as _MaxPool1d, MaxPool2d as _MaxPool2d
 
 from refiners.fluxion.layers.module import Module
 
 
-class MaxPool1d(nn.MaxPool1d, Module):
+class MaxPool1d(_MaxPool1d, Module):
+    """MaxPool1d layer.
+
+    This layer wraps [`torch.nn.MaxPool1d`][torch.nn.MaxPool1d].
+
+    Receives:
+        (Float[Tensor, "batch channels in_length"]):
+
+    Returns:
+        (Float[Tensor, "batch channels out_length"]):
+    """
+
     def __init__(
         self,
         kernel_size: int,
@@ -13,6 +24,16 @@ class MaxPool1d(nn.MaxPool1d, Module):
         return_indices: bool = False,
         ceil_mode: bool = False,
     ) -> None:
+        """Initializes the MaxPool1d layer.
+
+        Args:
+            kernel_size: The size of the sliding window.
+            stride: The stride of the sliding window.
+            padding: The amount of zero-padding added to both sides of the input.
+            dilation: The spacing between kernel elements.
+            return_indices: If True, returns the max indices along with the outputs.
+            ceil_mode: If True, uses ceil instead of floor to compute the output shape.
+        """
         super().__init__(
             kernel_size=kernel_size,
             stride=stride,
@@ -23,7 +44,18 @@ class MaxPool1d(nn.MaxPool1d, Module):
         )
 
 
-class MaxPool2d(nn.MaxPool2d, Module):
+class MaxPool2d(_MaxPool2d, Module):
+    """MaxPool2d layer.
+
+    This layer wraps [`torch.nn.MaxPool2d`][torch.nn.MaxPool2d].
+
+    Receives:
+        (Float[Tensor, "batch channels in_height in_width"]):
+
+    Returns:
+        (Float[Tensor, "batch channels out_height out_width"]):
+    """
+
     def __init__(
         self,
         kernel_size: int | tuple[int, int],
@@ -33,10 +65,20 @@ class MaxPool2d(nn.MaxPool2d, Module):
         return_indices: bool = False,
         ceil_mode: bool = False,
     ) -> None:
+        """Initializes the MaxPool2d layer.
+
+        Args:
+            kernel_size: The size of the sliding window.
+            stride: The stride of the sliding window.
+            padding: The amount of zero-padding added to both sides of the input.
+            dilation: The spacing between kernel elements.
+            return_indices: If True, returns the max indices along with the outputs.
+            ceil_mode: If True, uses ceil instead of floor to compute the output shape.
+        """
         super().__init__(
             kernel_size=kernel_size,
             stride=stride,
-            padding=padding,  # type: ignore
+            padding=padding,
             dilation=dilation,
             return_indices=return_indices,
             ceil_mode=ceil_mode,