← back to Exo
.typings/mlx/nn/layers/pooling.pyi
243 lines
"""
This type stub file was generated by pyright.
"""
from typing import Optional, Tuple, Union
import mlx.core as mx
from base import Module
class _Pool(Module):
def __init__(
self, pooling_function, kernel_size, stride, padding, padding_value
) -> None: ...
def __call__(self, x: mx.array) -> mx.array: ...
class _Pool1d(_Pool):
def __init__(
self,
pooling_function,
padding_value,
kernel_size: Union[int, Tuple[int]],
stride: Optional[Union[int, Tuple[int]]] = ...,
padding: Union[int, Tuple[int]] = ...,
) -> None: ...
class _Pool2d(_Pool):
def __init__(
self,
pooling_function,
padding_value,
kernel_size: Union[int, Tuple[int, int]],
stride: Optional[Union[int, Tuple[int, int]]] = ...,
padding: Optional[Union[int, Tuple[int, int]]] = ...,
) -> None: ...
class _Pool3d(_Pool):
def __init__(
self,
pooling_function,
padding_value,
kernel_size: Union[int, Tuple[int, int, int]],
stride: Optional[Union[int, Tuple[int, int, int]]] = ...,
padding: Optional[Union[int, Tuple[int, int, int]]] = ...,
) -> None: ...
class MaxPool1d(_Pool1d):
r"""Applies 1-dimensional max pooling.
Spatially downsamples the input by taking the maximum of a sliding window
of size ``kernel_size`` and sliding stride ``stride``.
Args:
kernel_size (int or tuple(int)): The size of the pooling window kernel.
stride (int or tuple(int), optional): The stride of the pooling window.
Default: ``kernel_size``.
padding (int or tuple(int), optional): How much negative infinity
padding to apply to the input. The padding amount is applied to
both sides of the spatial axis. Default: ``0``.
Examples:
>>> import mlx.core as mx
>>> import layers as nn
>>> x = mx.random.normal(shape=(4, 16, 5))
>>> pool = nn.MaxPool1d(kernel_size=2, stride=2)
>>> pool(x)
"""
def __init__(
self,
kernel_size: Union[int, Tuple[int]],
stride: Optional[Union[int, Tuple[int]]] = ...,
padding: Union[int, Tuple[int]] = ...,
) -> None: ...
class AvgPool1d(_Pool1d):
r"""Applies 1-dimensional average pooling.
Spatially downsamples the input by taking the average of a sliding window
of size ``kernel_size`` and sliding stride ``stride``.
Args:
kernel_size (int or tuple(int)): The size of the pooling window kernel.
stride (int or tuple(int), optional): The stride of the pooling window.
Default: ``kernel_size``.
padding (int or tuple(int), optional): How much zero padding to apply to
the input. The padding amount is applied to both sides of the spatial
axis. Default: ``0``.
Examples:
>>> import mlx.core as mx
>>> import layers as nn
>>> x = mx.random.normal(shape=(4, 16, 5))
>>> pool = nn.AvgPool1d(kernel_size=2, stride=2)
>>> pool(x)
"""
def __init__(
self,
kernel_size: Union[int, Tuple[int]],
stride: Optional[Union[int, Tuple[int]]] = ...,
padding: Union[int, Tuple[int]] = ...,
) -> None: ...
class MaxPool2d(_Pool2d):
r"""Applies 2-dimensional max pooling.
Spatially downsamples the input by taking the maximum of a sliding window
of size ``kernel_size`` and sliding stride ``stride``.
The parameters ``kernel_size``, ``stride``, and ``padding`` can either be:
* a single ``int`` -- in which case the same value is used for both the
height and width axis.
* a ``tuple`` of two ``int`` s -- in which case, the first ``int`` is
used for the height axis, the second ``int`` for the width axis.
Args:
kernel_size (int or tuple(int, int)): The size of the pooling window.
stride (int or tuple(int, int), optional): The stride of the pooling
window. Default: ``kernel_size``.
padding (int or tuple(int, int), optional): How much negative infinity
padding to apply to the input. The padding is applied on both sides
of the height and width axis. Default: ``0``.
Examples:
>>> import mlx.core as mx
>>> import layers as nn
>>> x = mx.random.normal(shape=(8, 32, 32, 4))
>>> pool = nn.MaxPool2d(kernel_size=2, stride=2)
>>> pool(x)
"""
def __init__(
self,
kernel_size: Union[int, Tuple[int, int]],
stride: Optional[Union[int, Tuple[int, int]]] = ...,
padding: Optional[Union[int, Tuple[int, int]]] = ...,
) -> None: ...
class AvgPool2d(_Pool2d):
r"""Applies 2-dimensional average pooling.
Spatially downsamples the input by taking the average of a sliding window
of size ``kernel_size`` and sliding stride ``stride``.
The parameters ``kernel_size``, ``stride``, and ``padding`` can either be:
* a single ``int`` -- in which case the same value is used for both the
height and width axis.
* a ``tuple`` of two ``int`` s -- in which case, the first ``int`` is
used for the height axis, the second ``int`` for the width axis.
Args:
kernel_size (int or tuple(int, int)): The size of the pooling window.
stride (int or tuple(int, int), optional): The stride of the pooling
window. Default: ``kernel_size``.
padding (int or tuple(int, int), optional): How much zero
padding to apply to the input. The padding is applied on both sides
of the height and width axis. Default: ``0``.
Examples:
>>> import mlx.core as mx
>>> import layers as nn
>>> x = mx.random.normal(shape=(8, 32, 32, 4))
>>> pool = nn.AvgPool2d(kernel_size=2, stride=2)
>>> pool(x)
"""
def __init__(
self,
kernel_size: Union[int, Tuple[int, int]],
stride: Optional[Union[int, Tuple[int, int]]] = ...,
padding: Optional[Union[int, Tuple[int, int]]] = ...,
) -> None: ...
class MaxPool3d(_Pool3d):
r"""Applies 3-dimensional max pooling.
Spatially downsamples the input by taking the maximum of a sliding window
of size ``kernel_size`` and sliding stride ``stride``.
The parameters ``kernel_size``, ``stride``, and ``padding`` can either be:
* a single ``int`` -- in which case the same value is used for the depth,
height, and width axis.
* a ``tuple`` of three ``int`` s -- in which case, the first ``int`` is used
for the depth axis, the second ``int`` for the height axis, and the third
``int`` for the width axis.
Args:
kernel_size (int or tuple(int, int, int)): The size of the pooling window.
stride (int or tuple(int, int, int), optional): The stride of the pooling
window. Default: ``kernel_size``.
padding (int or tuple(int, int, int), optional): How much negative infinity
padding to apply to the input. The padding is applied on both sides
of the depth, height and width axis. Default: ``0``.
Examples:
>>> import mlx.core as mx
>>> import layers as nn
>>> x = mx.random.normal(shape=(8, 16, 32, 32, 4))
>>> pool = nn.MaxPool3d(kernel_size=2, stride=2)
>>> pool(x)
"""
def __init__(
self,
kernel_size: Union[int, Tuple[int, int, int]],
stride: Optional[Union[int, Tuple[int, int, int]]] = ...,
padding: Optional[Union[int, Tuple[int, int, int]]] = ...,
) -> None: ...
class AvgPool3d(_Pool3d):
r"""Applies 3-dimensional average pooling.
Spatially downsamples the input by taking the average of a sliding window
of size ``kernel_size`` and sliding stride ``stride``.
The parameters ``kernel_size``, ``stride``, and ``padding`` can either be:
* a single ``int`` -- in which case the same value is used for the depth,
height, and width axis.
* a ``tuple`` of three ``int`` s -- in which case, the first ``int`` is used
for the depth axis, the second ``int`` for the height axis, and the third
``int`` for the width axis.
Args:
kernel_size (int or tuple(int, int, int)): The size of the pooling window.
stride (int or tuple(int, int, int), optional): The stride of the pooling
window. Default: ``kernel_size``.
padding (int or tuple(int, int, int), optional): How much zero
padding to apply to the input. The padding is applied on both sides
of the depth, height and width axis. Default: ``0``.
Examples:
>>> import mlx.core as mx
>>> import layers as nn
>>> x = mx.random.normal(shape=(8, 16, 32, 32, 4))
>>> pool = nn.AvgPool3d(kernel_size=2, stride=2)
>>> pool(x)
"""
def __init__(
self,
kernel_size: Union[int, Tuple[int, int, int]],
stride: Optional[Union[int, Tuple[int, int, int]]] = ...,
padding: Optional[Union[int, Tuple[int, int, int]]] = ...,
) -> None: ...