exo/.mlx_typings/mlx/nn/layers/pooling.pyi

"""
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: ...