# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.

from .module import Module
from ... import Tensor
from ..common_types import _size_any_t


class Fold(Module):
    output_size: _size_any_t = ...
    kernel_size: _size_any_t = ...
    dilation: _size_any_t = ...
    padding: _size_any_t = ...
    stride: _size_any_t = ...

    def __init__(self, output_size: _size_any_t, kernel_size: _size_any_t, dilation: _size_any_t = ...,
                 padding: _size_any_t = ..., stride: _size_any_t = ...) -> None: ...

    def forward(self, input: Tensor) -> Tensor: ...  # type: ignore

    def __call__(self, input: Tensor) -> Tensor: ...  # type: ignore


class Unfold(Module):
    kernel_size: _size_any_t = ...
    dilation: _size_any_t = ...
    padding: _size_any_t = ...
    stride: _size_any_t = ...

    def __init__(self, kernel_size: _size_any_t, dilation: _size_any_t = ..., padding: _size_any_t = ...,
                 stride: _size_any_t = ...) -> None: ...

    def forward(self, input: Tensor) -> Tensor: ...  # type: ignore

    def __call__(self, input: Tensor) -> Tensor: ...  # type: ignore