`.typings/mlx/nn/layers/convolution_transpose.pyi`

120 lines
"""
This type stub file was generated by pyright.
"""

from typing import Union

import mlx.core as mx
from base import Module

class ConvTranspose1d(Module):
    """Applies a 1-dimensional transposed convolution over the multi-channel input sequence.

    The channels are expected to be last i.e. the input shape should be ``NLC`` where:

    * ``N`` is the batch dimension
    * ``L`` is the sequence length
    * ``C`` is the number of input channels

    Args:
        in_channels (int): The number of input channels
        out_channels (int): The number of output channels
        kernel_size (int): The size of the convolution filters
        stride (int, optional): The stride when applying the filter.
            Default: ``1``.
        padding (int, optional): How many positions to 0-pad the input with.
            Default: ``0``.
        dilation (int, optional): The dilation of the convolution.
        output_padding(int, optional): Additional size added to one side of the
            output shape. Default: ``0``.
        bias (bool, optional): If ``True`` add a learnable bias to the output.
            Default: ``True``
    """
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: int,
        stride: int = ...,
        padding: int = ...,
        dilation: int = ...,
        output_padding: int = ...,
        bias: bool = ...,
    ) -> None: ...
    def __call__(self, x: mx.array) -> mx.array: ...

class ConvTranspose2d(Module):
    """Applies a 2-dimensional transposed convolution over the multi-channel input image.

    The channels are expected to be last i.e. the input shape should be ``NHWC`` where:

    * ``N`` is the batch dimension
    * ``H`` is the input image height
    * ``W`` is the input image width
    * ``C`` is the number of input channels

    Args:
        in_channels (int): The number of input channels.
        out_channels (int): The number of output channels.
        kernel_size (int or tuple): The size of the convolution filters.
        stride (int or tuple, optional): The size of the stride when
            applying the filter. Default: ``1``.
        padding (int or tuple, optional): How many positions to 0-pad
            the input with. Default: ``0``.
        dilation (int or tuple, optional): The dilation of the convolution.
        output_padding(int or tuple, optional): Additional size added to one
            side of the output shape. Default: ``0``.
        bias (bool, optional): If ``True`` add a learnable bias to the
            output. Default: ``True``
    """
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: Union[int, tuple],
        stride: Union[int, tuple] = ...,
        padding: Union[int, tuple] = ...,
        dilation: Union[int, tuple] = ...,
        output_padding: Union[int, tuple] = ...,
        bias: bool = ...,
    ) -> None: ...
    def __call__(self, x: mx.array) -> mx.array: ...

class ConvTranspose3d(Module):
    """Applies a 3-dimensional transposed convolution over the multi-channel input image.

    The channels are expected to be last i.e. the input shape should be ``NDHWC`` where:

    * ``N`` is the batch dimension
    * ``D`` is the input image depth
    * ``H`` is the input image height
    * ``W`` is the input image width
    * ``C`` is the number of input channels

    Args:
        in_channels (int): The number of input channels.
        out_channels (int): The number of output channels.
        kernel_size (int or tuple): The size of the convolution filters.
        stride (int or tuple, optional): The size of the stride when
            applying the filter. Default: ``1``.
        padding (int or tuple, optional): How many positions to 0-pad
            the input with. Default: ``0``.
        dilation (int or tuple, optional): The dilation of the convolution.
        output_padding(int or tuple, optional): Additional size added to one
            side of the output shape. Default: ``0``.
        bias (bool, optional): If ``True`` add a learnable bias to the
            output. Default: ``True``
    """
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        kernel_size: Union[int, tuple],
        stride: Union[int, tuple] = ...,
        padding: Union[int, tuple] = ...,
        dilation: Union[int, tuple] = ...,
        output_padding: Union[int, tuple] = ...,
        bias: bool = ...,
    ) -> None: ...
    def __call__(self, x: mx.array) -> mx.array: ...