from collections.abc import Callable, Iterator
from pandas import DataFrame
from torch import Tensor, dtype, device
from torch.nn import Module, Dropout, AlphaDropout
from torch.optim import Optimizer
from torch.optim.lr_scheduler import LRScheduler

type Functional = Callable[[Tensor], Tensor]
type Tensors1T = tuple[Tensor]
type Tensors2T = tuple[Tensor, Tensor]
type Tensors3T = tuple[Tensor, Tensor, Tensor]
type Tensors4T = tuple[Tensor, Tensor, Tensor, Tensor]
type Tensors = tuple[Tensor, ...]
type Dtype = dtype
type Device = device
type Drop = Dropout | AlphaDropout
type Batch = tuple[Tensors, Tensor]
type Batches = Iterator[Batch]


__all__ = [
    'DataFrame',
    'Tensor',
    'Module',
    'Functional',
    'Tensors1T',
    'Tensors2T',
    'Tensors3T',
    'Tensors4T',
    'Tensors',
    'Dtype',
    'Device',
    'Drop',
    'Optimizer',
    'LRScheduler',
    'Batch',
    'Batches'
]
