(doc/fluxion/maxpool) add/convert docstrings to mkdocstrings format

This commit is contained in:
Laurent 2024-02-01 21:46:23 +00:00 committed by Laureηt
parent 49847658e9
commit 18682f8155

View file

@ -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,