mirror of
https://github.com/finegrain-ai/refiners.git
synced 2024-11-22 14:18:46 +00:00
(doc/fluxion/maxpool) add/convert docstrings to mkdocstrings format
This commit is contained in:
parent
49847658e9
commit
18682f8155
|
@ -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
|
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__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
kernel_size: int,
|
kernel_size: int,
|
||||||
|
@ -13,6 +24,16 @@ class MaxPool1d(nn.MaxPool1d, Module):
|
||||||
return_indices: bool = False,
|
return_indices: bool = False,
|
||||||
ceil_mode: bool = False,
|
ceil_mode: bool = False,
|
||||||
) -> None:
|
) -> 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__(
|
super().__init__(
|
||||||
kernel_size=kernel_size,
|
kernel_size=kernel_size,
|
||||||
stride=stride,
|
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__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
kernel_size: int | tuple[int, int],
|
kernel_size: int | tuple[int, int],
|
||||||
|
@ -33,10 +65,20 @@ class MaxPool2d(nn.MaxPool2d, Module):
|
||||||
return_indices: bool = False,
|
return_indices: bool = False,
|
||||||
ceil_mode: bool = False,
|
ceil_mode: bool = False,
|
||||||
) -> None:
|
) -> 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__(
|
super().__init__(
|
||||||
kernel_size=kernel_size,
|
kernel_size=kernel_size,
|
||||||
stride=stride,
|
stride=stride,
|
||||||
padding=padding, # type: ignore
|
padding=padding,
|
||||||
dilation=dilation,
|
dilation=dilation,
|
||||||
return_indices=return_indices,
|
return_indices=return_indices,
|
||||||
ceil_mode=ceil_mode,
|
ceil_mode=ceil_mode,
|
||||||
|
|
Loading…
Reference in a new issue