# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/03_kwds.ipynb.

# %% auto 0
__all__ = ['NeuralKeywords', 'AugmenterKeywords', 'SimpleKeywords', 'DynamicsForwardKeywords', 'DynamicsKeywords',
           'NeuralOptimalTransportKeywords', 'DifferentialEquationKeywords', 'ODEKeywords', 'SDEFunctionKeywords',
           'SDEProblemKeywords', 'SDEKeywords', 'SubVidInfo', 'VidInfo', 'DiffusionKeywords', 'BaseRecurrentKeywords',
           'LSTMKeywords', 'RNNKeywords', 'RecurrentKeywords', 'RecurrentAutoEncoderKeywords',
           'RecurrentEncoderKeywords', 'RecurrentDecoderKeywords']

# %% ../nbs/03_kwds.ipynb 6
from functools import wraps

# %% ../nbs/03_kwds.ipynb 8
from typing import (
    Union, Iterable, TypedDict, 
    Optional, Callable, Generator, overload, get_overloads
)

# %% ../nbs/03_kwds.ipynb 10
#| export


# %% ../nbs/03_kwds.ipynb 12
from nlit import KIND, NONLINEARITY, PROJ_SIZE, INPUT_SIZE, HIDDEN_SIZE, OUTPUT_SIZE
from quac import (
    intq, tensor, deviceq, dtypeq, xypos, rect, path,
    nnmodule, neuraldynamics, nnmodule, sdefunc, brownianintervial, bbox as BBox
)
from dtyp import TypeDict, new
from etrc import Solver, SolverKind, SDENoiseType, NonLinearity, RecurrentLayer, DynamicsMethod

# %% ../nbs/03_kwds.ipynb 14
#| export


# %% ../nbs/03_kwds.ipynb 16
_IGNORE_INIT = dict(
    tensor = tensor, 
    nnmodule = nnmodule,
)

_TYPE_DICT_KWDS = dict(
    __ignore_values  = _IGNORE_INIT,
    __force_defaults = False
)

# %% ../nbs/03_kwds.ipynb 19
class NeuralKeywords(TypeDict, **_TYPE_DICT_KWDS):
    '''Typed dictionary for initalizing a neural network.
    
    Attributes
    ----------
    device : device | str, optional
        The device to use for computation.
    '''
    device: deviceq # device to use for computation
    
    @overload
    def __init__(self, device: deviceq) -> 'NeuralKeywords': ...
    def __init__(self, *args, **kwargs: 'NeuralKeywords') -> 'NeuralKeywords':
        '''Typed dictionary for initalizing a neural network.'''
        return super().__init__(self, *(), **kwargs)

# %% ../nbs/03_kwds.ipynb 22
class AugmenterKeywords(TypeDict, **_TYPE_DICT_KWDS):
    '''Typed dictionary for initalizing an Augmenter.
    
    Attributes
    ----------
    adim : int, optional
        The number of augmenting dimensions.
        
    afnc : Callable, optional
        The augmenting function.
    '''
    adim: intq # number of augmenting dimensions
    afnc: Callable | None # augment_func
    
    @overload
    def __init__(self, adim: intq, afnc: Callable | None) -> 'AugmenterKeywords': ...
    def __init__(self, *args, **kwargs: 'AugmenterKeywords') -> 'AugmenterKeywords':
        '''Typed dictionary for initalizing an Augmenter.'''
        return super().__init__(self, *(), **kwargs)

# %% ../nbs/03_kwds.ipynb 25
class SimpleKeywords(NeuralKeywords, **_TYPE_DICT_KWDS):
    '''Typed dictionary for initalizing a Simple Neural Network.
    
    Attributes
    ----------
    feats : list[int]
        The number of features per linear layer e.g. [10, 10]
    
    act : NonLinearity
        The activation function to use for linear layers.
    
    endact : bool
        Whether to apply activation function to last linear layer.
        
    out : NonLinearity
        The activation function to use for the output e.g. Softmax.
        
    outact : bool
        Whether to apply activation function to the output
        (e.g. after layer[0] --> (act) --> ... --> layer[-1] --> (act?<endact>) --> (finact?<outact>))
        
    outkws : dict
        Keyword arguments to pass to the output activation function.
        
    device : device | str, optional
        The device to use for computation.
    '''
    
    feats: list[int]  # number of features per linear layer e.g. [10, 10]
    act: NonLinearity = NonLinearity.CELU # activation function to use for linear layers
    endact: bool = False # whether to apply activation function to last linear layer
    out: NonLinearity = NonLinearity.Softmax # activation function to use for the output e.g. Softmax
    outact: bool = False # whether to apply activation function to the output
                         # (e.g. after layer[0] --> (act) --> ... --> layer[-1] --> (act?<endact>) --> (finact?<outact>))
    outkws: dict = dict(dim=1) # keyword arguments to pass to the output activation function
    @overload
    def __init__(
        self, feats: list[int], 
        act: NonLinearity, 
        endact: bool, 
        out: NonLinearity, 
        outact: bool, 
        outkws: dict,
        **kwargs: NeuralKeywords
    ) -> 'SimpleKeywords': ...
    @overload
    def __init__(
        self, feats: list[int], 
        act: NonLinearity, 
        endact: bool, 
        out: NonLinearity, 
        outact: bool, 
        outkws: dict,
        device: deviceq
    ) -> 'SimpleKeywords': ...
    def __init__(self, *args, **kwargs: 'SimpleKeywords') -> 'SimpleKeywords':
        '''Typed dictionary for initalizing a Simple Neural Network.'''
        return super().__init__(self, *(), **kwargs)

# %% ../nbs/03_kwds.ipynb 29
class DynamicsForwardKeywords(TypeDict, **_TYPE_DICT_KWDS):
    '''Typed dictionary for the forward call of a `torchdyn.nn.Module`.
    
    Attributes
    ----------
    x : tensor
    
    t : tensor, optional
    
    save_at : Iterable
    
    args : dict
    
    aug : nnmodule, optional
    
    dyn : neuraldynamics, optional
    
    dec : nnmodule, optional
    
    deltas : list[int], optional
    
    useaug : bool
    
    usedec : bool
    '''
    x: tensor
    t: Optional[tensor]
    save_at: Iterable = ()
    args: dict = dict()
    aug: Optional[nnmodule]
    dyn: Optional[neuraldynamics]
    dec: Optional[nnmodule]
    deltas: Optional[list[int]]
    useaug: bool = True
    usedec: bool = True
    
    @overload
    def __init__(
        self, x: tensor,
        t: Optional[tensor],
        save_at: Iterable,
        args: dict,
        aug: Optional[nnmodule],
        dyn: Optional[neuraldynamics],
        dec: Optional[nnmodule],
        deltas: Optional[list[int]],
        useaug: bool, 
        usedec: bool,
    ) -> 'DynamicsForwardKeywords': ...
    def __init__(self, *args, **kwargs: 'DynamicsForwardKeywords') -> 'DynamicsForwardKeywords':
        '''Typed dictionary for the forward call of a `torchdyn.nn.Module`.'''
        return super().__init__(self, *(), **kwargs)

# %% ../nbs/03_kwds.ipynb 32
class DynamicsKeywords(AugmenterKeywords, NeuralKeywords, **_TYPE_DICT_KWDS):
    '''Typed dictionary for initalizing a Dynamic Neural Network.
    
    Attributes
    ----------
    diffeq : DynamicsMethod 
        The dynamics method to use e.g. DynamicsMethod.NeuralODE, DynamicsMethod.NeuralSDE
        
    tsteps : tensor
        The time steps to consider e.g. [0, 1, 2, 3]
    
    adim : int, optional
        The number of augmenting dimensions.
        
    afnc : Callable, optional
        The augmenting function.
        
    device : device | str, optional
        The device to use for computation.
    '''
    diffeq: DynamicsMethod = DynamicsMethod.NeuralODE # dynamics method to use e.g. DynamicsMethod.NeuralODE, DynamicsMethod.NeuralSDE
    tsteps: tensor         # time steps to consider e.g. [0, 1, 2, 3]
    
    @overload
    def __init__(
        self, 
        diffeq: DynamicsMethod,
        tsteps: tensor,
        adim: intq, afnc: Callable | None,
        device : deviceq
    ): ...
    @overload
    def __init__(
        self, 
        diffeq: DynamicsMethod,
        tsteps: tensor,
        adim: intq, afnc: Callable | None,
        **kwargs: NeuralKeywords
    ): ...
    @overload
    def __init__(
        self, 
        diffeq: DynamicsMethod,
        tsteps: tensor,
        **kwargs: AugmenterKeywords | NeuralKeywords
    ): ...
    def __init__(self, *args, **kwargs: 'DynamicsForwardKeywords') -> 'DynamicsForwardKeywords':
        '''Typed dictionary for initalizing a Dynamic Neural Network.'''
        return super().__init__(self, *(), **kwargs)

# %% ../nbs/03_kwds.ipynb 35
class NeuralOptimalTransportKeywords(SimpleKeywords, DynamicsKeywords, AugmenterKeywords, NeuralKeywords, **_TYPE_DICT_KWDS):
    '''Typed dictionary for initalizing an Optimal Transport Neural Network.
    
    Attributes
    ----------
    topk: int
        number of top predictions to consider
        
    feats: list[int]
        The number of features per linear layer e.g. [10, 10]
    
    act: NonLinearity
        The activation function to use for linear layers.
    
    endact: bool
        Whether to apply activation function to last linear layer.
        
    out: NonLinearity
        The activation function to use for the output e.g. Softmax.
        
    outact: bool
        Whether to apply activation function to the output
        (e.g. after layer[0] --> (act) --> ... --> layer[-1] --> (act?<endact>) --> (finact?<outact>))
        
    outkws: dict
        Keyword arguments to pass to the output activation function.
        
    diffeq: DynamicsMethod 
        The dynamics method to use e.g. DynamicsMethod.NeuralODE, DynamicsMethod.NeuralSDE
        
    tsteps: tensor
        The time steps to consider e.g. [0, 1, 2, 3]
    
    adim : int | None
        The number of augmenting dimensions.
        
    afnc : Callable | None
        The augmenting function.
        
    device : device | str | None
        The device to use for computation.
    '''
    topk: int = 5# number of top predictions to consider
    @overload
    def __init__(
        self, 
        diffeq: DynamicsMethod,
        tsteps: tensor,
        adim: intq, afnc: Callable | None,
        device : deviceq
    ): ...
    @overload
    def __init__(
        self, 
        diffeq: DynamicsMethod,
        tsteps: tensor,
        adim: intq, afnc: Callable | None,
        **kwargs: NeuralKeywords
    ): ...
    @overload
    def __init__(
        self, 
        diffeq: DynamicsMethod,
        tsteps: tensor,
        **kwargs: AugmenterKeywords | NeuralKeywords
    ): ...
    @overload
    def __init__(
        self, 
        topk: int,
        **kwargs: DynamicsKeywords | AugmenterKeywords | NeuralKeywords
    ): ...
    def __init__(self, *args, **kwargs: 'NeuralOptimalTransportKeywords') -> 'NeuralOptimalTransportKeywords':
        '''Typed dictionary for initalizing an Optimal Transport Neural Network.'''
        return super().__init__(self, *(), **kwargs)

# %% ../nbs/03_kwds.ipynb 38
class DifferentialEquationKeywords(TypeDict, **_TYPE_DICT_KWDS):
    '''Typed dictionary for initalizing a Differential Equation.
    
    Attributes
    ----------
    order: Optional[int], default: 1
        order of the differential equation. Defaults to 1.
        
    solver: Union[Solver, str, nnmodule], default: Solver.EulerODE
        Defaults to Solver.EulerODE.
    
    atol: float, default: 1e-4
        Defaults to 1e-4.
        
    rtol: float, default: 1e-4
        Defaults to 1e-4.
        
    sensitivity: str, default: 'autograd'
        Defaults to 'autograd'.
        
    return_t_eval: bool, default: True
        Defaults to True.
    '''
    order: Optional[int] = 1#  Defaults to 1. i.e. order (int, optional): order of the differential equation. Defaults to 1.    
    solver: Union[Solver, str, nnmodule] = Solver.EulerODE.value.name.lower() # Defaults to 'euler'.
    atol: float = 1e-4 # Defaults to 1e-4.
    rtol: float = 1e-4 # Defaults to 1e-4.
    sensitivity: str = 'autograd' # Defaults to 'autograd'.
    return_t_eval: bool = True # Defaults to True.
    
    @overload
    def __init__(
        self, 
        order: intq,
        solver: Union[Solver, str, nnmodule],
        atol: float,
        rtol: float,
        sensitivity: str,
        return_t_eval: bool,
    ): ...
    def __init__(self, *args, **kwargs: 'DifferentialEquationKeywords') -> 'DifferentialEquationKeywords':
        '''Typed dictionary for initalizing a Differential Equation.'''
        return super().__init__(self, *(), **kwargs)

# %% ../nbs/03_kwds.ipynb 42
class ODEKeywords(DifferentialEquationKeywords, **_TYPE_DICT_KWDS):
    '''Typed dictionary for initalizing an Ordinary Differential Equation.
    
    Attributes
    ----------
    vector_field: Union[Solver, str, nnmodule]
        Defaults to Solver.EulerODE.
    
    solver_adjoint: Optional[Union[Solver, str, nnmodule]]
        Defaults to None.
        
    atol_adjoint: float
        Defaults to 1e-4.
        
    rtol_adjoint: float
        Defaults to 1e-4.
        
    interpolator: Union[str, Callable, None]
        Defaults to None.
        
    integral_loss: Union[Callable, None]
        Defaults to None.
        
    seminorm: bool
        Defaults to False.
        
    optimizable_params: Union[Iterable, Generator]
        Defaults to ().
    
    order: Optional[int]
        order of the differential equation. Defaults to 1.
        
    solver: Union[Solver, str, nnmodule]
        Defaults to Solver.EulerODE.
    
    atol: float
        Defaults to 1e-4.
        
    rtol: float
        Defaults to 1e-4.
        
    sensitivity: str
        Defaults to 'autograd'.
        
    return_t_eval: bool
        Defaults to True.
    '''
    vector_field: Union[Solver, str, nnmodule] = Solver.EulerODE.value.name.lower()
    solver_adjoint: Optional[Union[Solver, str, nnmodule]] = None
    atol_adjoint: float = 1e-4 # defaults to 1e-4.
    rtol_adjoint: float = 1e-4 # defaults to 1e-4.
    interpolator: Union[str, Callable, None] = None
    integral_loss: Union[Callable, None] = None
    seminorm: bool = False
    optimizable_params: Union[Iterable, Generator] = (),
    
    @overload
    def __init__(
        self, 
        order: intq,
        solver: Union[Solver, str, nnmodule],
        atol: float,
        rtol: float,
        sensitivity: str,
        return_t_eval: bool,
    ): ...
    @overload
    def __init__(
        self, 
        vector_field: Union[Solver, str, nnmodule],
        solver_adjoint: Optional[Union[Solver, str, nnmodule]],
        atol_adjoint: float,
        rtol_adjoint: float,
        interpolator: Union[str, Callable, None],
        integral_loss: Union[Callable, None],
        seminorm: bool,
        optimizable_params: Union[Iterable, Generator],
        **kwargs: DifferentialEquationKeywords
    ): ...
    def __init__(self, *args, **kwargs: 'ODEKeywords') -> 'ODEKeywords':
        '''Typed dictionary for initalizing an Ordinary Differential Equation.'''
        return super().__init__(self, *(), **kwargs)

# %% ../nbs/03_kwds.ipynb 46
class SDEFunctionKeywords(TypeDict, **_TYPE_DICT_KWDS):
    '''Typed dictionary for initalizing a SDE function.
    
    Attributes
    ----------
    drift_func: Callable
        drift function, i.e. f (Callable): callable defining the drift
    
    diffusion_func: Callable
        diffusion function i.e. g (Callable): callable defining the diffusion term
    
    order: Optional[int]
        order of the differential equation. Defaults to 1.
    
    noise_type: SDENoiseType
        Defaults to SDENoiseType.diagonal.
    
    sde_type: Optional[SolverKind]
        Defaults to SolverKind.ito.
    '''
    drift_func: Callable      # drift function, i.e. f (Callable): callable defining the drift
    diffusion_func: Callable  # diffusion function i.e. g (Callable): callable defining the diffusion term
    order: Optional[int] = 1  # Defaults to 1. i.e. order (int, optional): order of the differential equation. Defaults to 1.
    noise_type: SDENoiseType = SDENoiseType.diagonal # Defaults to 'SDENoiseType.diagonal'.
    sde_type: Optional[SolverKind] = SolverKind.ito # Defaults to 'SolverKind.ito'.
    @overload
    def __init__(
        self, 
        drift_func: Callable,
        diffusion_func: Callable,
        order: intq,
        noise_type: SDENoiseType,
        sde_type: Optional[SolverKind],
    ): ...
    def __init__(self, *args, **kwargs: 'SDEFunctionKeywords') -> 'SDEFunctionKeywords':
        '''Typed dictionary for initalizing a SDE function.'''
        return super().__init__(self, *(), **kwargs)

# %% ../nbs/03_kwds.ipynb 49
class SDEProblemKeywords(TypeDict, **_TYPE_DICT_KWDS):
    '''Typed dictionary for initalizing a SDE Problem function.
    
    Attributes
    ----------
    defunc: sdefunc
        The differnetial equation function e.g. 
        `SDEFunc(f=drift, g=diffusion, order=order)`
    '''
    defunc: sdefunc # NOTE: defunc <== SDEFunc(f=drift_func, g=diffusion_func, order=order)
    @overload
    def __init__(
        self, 
        defunc: sdefunc,
    ): ...
    def __init__(self, *args, **kwargs: 'SDEProblemKeywords') -> 'SDEProblemKeywords':
        '''Typed dictionary for initalizing a SDE Problem function.'''
        return super().__init__(self, *(), **kwargs)

# %% ../nbs/03_kwds.ipynb 52
class SDEKeywords(SDEFunctionKeywords, DifferentialEquationKeywords, **_TYPE_DICT_KWDS):
    '''Typed dictionary for initalizing a Stochastic Differential Equation.
    
    Attributes
    ----------
    t_span: Optional[list[int]]
        Defaults to torch.linspace(0, 1, 2).
        
    ds: float
        Defaults to 1e-3.
        
    intloss: Optional[int]
        Defaults to None.
        
    bm: Optional[brownianintervial]
        Defaults to Brownian Motion.
    
    drift_func: Callable
        drift function, i.e. f (Callable): callable defining the drift
    
    diffusion_func: Callable
        diffusion function i.e. g (Callable): callable defining the diffusion term
    
    order: Optional[CatOrder]
        order of the differential equation. Defaults to CatOrder.first.
    
    noise_type: SDENoiseType
        Defaults to SDENoiseType.diagonal.
    
    sde_type: Optional[SolverKind]
        Defaults to SolverKind.ito.
        
    vector_field: Union[Solver, str, nnmodule]
        Defaults to Solver.EulerODE.
    
    solver_adjoint: Optional[Union[Solver, str, nnmodule]]
        Defaults to None.
        
    atol_adjoint: float
        Defaults to 1e-4.
        
    rtol_adjoint: float
        Defaults to 1e-4.
        
    interpolator: Union[str, Callable, None]
        Defaults to None.
        
    integral_loss: Union[Callable, None]
        Defaults to None.
        
    seminorm: bool
        Defaults to False.
        
    optimizable_params: Union[Iterable, Generator]
        Defaults to ().
    
    order: Optional[CatOrder]
        order of the differential equation. Defaults to CatOrder.first.
        
    solver: Union[Solver, str, nnmodule]
        Defaults to Solver.EulerODE.
    
    atol: float
        Defaults to 1e-4.
        
    rtol: float
        Defaults to 1e-4.
        
    sensitivity: str
        Defaults to 'autograd'.
        
    return_t_eval: bool
        Defaults to True.
    '''
    t_span: Optional[list[int]] # Defaults to torch.linspace(0, 1, 2).
    ds: float = 1e-3 # Defaults to 1e-3.
    intloss: intq = None # Defaults to None.
    bm: Optional[brownianintervial] # Brownian Motion
    @overload
    def __init__(
        self, 
        t_span: Optional[list[int]],
        ds: float,
        intloss: intq,
        bm: Optional[brownianintervial],
        drift_func: Callable,
        diffusion_func: Callable,
        order: intq,
        noise_type: SDENoiseType,
        sde_type: Optional[SolverKind],
        solver: Union[Solver, str, nnmodule],
        atol: float,
        rtol: float,
        sensitivity: str,
        return_t_eval: bool,
    ): ...
    @overload
    def __init__(
        self, 
        t_span: Optional[list[int]],
        ds: float,
        intloss: intq,
        bm: Optional[brownianintervial],
        drift_func: Callable,
        diffusion_func: Callable,
        order: intq,
        noise_type: SDENoiseType,
        sde_type: Optional[SolverKind],
        **kwargs: DifferentialEquationKeywords
    ): ...
    @overload
    def __init__(
        self, 
        t_span: Optional[list[int]],
        ds: float,
        intloss: intq,
        bm: Optional[brownianintervial],
        **kwargs: SDEFunctionKeywords | DifferentialEquationKeywords
    ): ...
    def __init__(self, *args, **kwargs: 'SDEKeywords') -> 'SDEKeywords':
        '''Typed dictionary for initalizing a Stochastic Differential Equation.'''
        return super().__init__(self, *(), **kwargs)

# %% ../nbs/03_kwds.ipynb 56
class SubVidInfo(TypeDict, **_TYPE_DICT_KWDS):
    '''Subvideo Information from an AVI file'''
    pos: xypos
    '''The (row, col) position of the subvideo in the grid'''
    bbox: BBox
    '''The bounding box of the subvideo in the grid'''
    subidx: int
    '''The flattened grid index of the subvideo'''
    subkey: str
    '''The label the subvideo'''
    @overload
    def __init__(
        self, 
        pos: xypos,
        bbox: BBox,
        subidx: int,
        subkey: str
    ): ...
    def __init__(self, *args, **kwargs: 'SubVidInfo') -> 'SubVidInfo':
        '''Typed dictionary containing Subvideo Information from an AVI file.'''
        return super().__init__(self, *(), **kwargs)

# %% ../nbs/03_kwds.ipynb 59
class VidInfo(TypeDict, **_TYPE_DICT_KWDS):
    '''AVI file information'''
    vidpxs: rect
    '''The frame size (width, height) of the video e.g. (1200, 1200)'''
    subpxs: rect
    '''The standardized subvideo size (sub-width, sub-height) e.g. (400, 400)'''
    values: list[dict[str, SubVidInfo]]
    '''The subvideo information'''
    source: path
    '''The source AVI file'''
    
    @overload
    def __init__(
        self, 
        vidpxs: rect, 
        subpxs: rect, 
        values: list[dict[str, SubVidInfo]], 
        source: path
    ): ...
    def __init__(self, *args, **kwargs: 'VidInfo') -> 'VidInfo':
        '''Typed dictionary containing AVI file information.'''
        return super().__init__(self, *(), **kwargs)

# %% ../nbs/03_kwds.ipynb 63
class DiffusionKeywords(TypeDict, **_TYPE_DICT_KWDS):
    '''Typed dictionary for initalizing a Diffusion Operation Method.'''
    knn: int = 5
    tmax: int = 5
    anisotropy: int = 1
    tdiff: int = 1
    topeig: int = 100
    log: bool = False
    normalize: bool = False
    symmetrize: bool = False
    verbose: bool = False
    nemb: int = 10
    njobs: int = -1
    device: deviceq
    
    @overload
    def __init__(
        self, 
        knn: int,
        tmax: int,
        anisotropy: int,
        tdiff: int,
        topeig: int,
        log: bool,
        normalize: bool,
        symmetrize: bool,
        verbose: bool,
        nemb: int,
        njobs: int,
        device: deviceq
    ): ...
    def __init__(self, *args, **kwargs: 'DiffusionKeywords') -> 'DiffusionKeywords':
        '''Typed dictionary for initalizing a Diffusion Operation Method.'''
        return super().__init__(self, *(), **kwargs)

# %% ../nbs/03_kwds.ipynb 67
class BaseRecurrentKeywords(NeuralKeywords, **_TYPE_DICT_KWDS):
    '''Typed dictionary for specifying base keywords for recurrent neural network layers.

    Attributes
    ----------
    input_size : int
        The number of expected features in the input.

    hidden_size : int
        The number of features in the hidden state.

    nlays : int
        Number of recurrent layers.

    bias : bool
        If False, then the layer does not use bias weights.

    batch_first : bool
        If True, then the input and output tensors are provided as (batch, seq, feature).

    dropout : float
        If non-zero, introduces a dropout layer on the outputs of each layer except the last layer.

    bidirectional : bool
        If True, becomes a bidirectional layer.

    device : deviceq
        The device on which to train the model.

    dtype : dtypeq
        The data type of the model parameters.

    Examples
    --------
    >>> base_rnn_keywords = BaseRecurrentKeywords(
    ...     input_size=64, hidden_size=128, nlays=2, bias=True, 
    ...     batch_first=True, dropout=0.5, bidirectional=False, device='cuda', 
    ...     dtype=torch.float32
    ... )
    '''
    input_size: int
    hidden_size: int
    nlays: int = 1
    
    bias: bool = True
    batch_first: bool = True
    dropout: float = 0.2
    bidirectional: bool = False
    
    device: deviceq
    dtype: dtypeq 
    
    @overload
    def __init__(
        self, 
        input_size: int,
        hidden_size: int,
        nlays: int,
        bias: bool,
        batch_first: bool ,
        dropout: float ,
        bidirectional: bool ,
        device: deviceq,
        dtype: dtypeq
    ): ...
    def __init__(self, *args, **kwargs: 'BaseRecurrentKeywords') -> 'BaseRecurrentKeywords':
        '''Typed dictionary for specifying base keywords for recurrent neural network layers.'''
        return super().__init__(self, *(), **kwargs)

# %% ../nbs/03_kwds.ipynb 71
class LSTMKeywords(BaseRecurrentKeywords, **_TYPE_DICT_KWDS):
    '''Typed dictionary for specifying keywords specific to LSTM layers, extending `BaseRecurrentKeywords`.

    Attributes
    ----------
    proj_size : int
        The size of the projection layer.

    Examples
    --------
    >>> lstm_keywords = LSTMKeywords(
    ...     input_size=64, hidden_size=128, nlays=2, proj_size=64, bias=True, 
    ...     batch_first=True, dropout=0.5, bidirectional=False, device='cuda', 
    ...     dtype=torch.float32
    ... )
    '''
    proj_size: int = 0
    
    @overload
    def __init__(
        self, 
        proj_size: int,
        input_size: int,
        hidden_size: int,
        nlays: int,
        bias: bool,
        batch_first: bool ,
        dropout: float ,
        bidirectional: bool ,
        device: deviceq,
        dtype: dtypeq,
    ): ...
    @overload
    def __init__(
        self, 
        proj_size: int,
        **kwargs: BaseRecurrentKeywords
    ): ...
    def __init__(self, *args, **kwargs: 'LSTMKeywords') -> 'LSTMKeywords':
        '''Typed dictionary for specifying keywords specific to LSTM layers, extending `BaseRecurrentKeywords`.'''
        return super().__init__(self, *(), **kwargs)

# %% ../nbs/03_kwds.ipynb 74
class RNNKeywords(BaseRecurrentKeywords, **_TYPE_DICT_KWDS):    
    '''Typed dictionary for specifying keywords specific to RNN layers, extending `BaseRecurrentKeywords`.

    Attributes
    ----------
    nonlinearity : NonLinearity
        The non-linearity to use ('relu' or 'tanh').

    Examples
    --------
    >>> rnn_keywords = RNNKeywords(
    ...     input_size=64, hidden_size=128, nlays=2, nonlinearity=NonLinearity.Tanh, 
    ...     bias=True, batch_first=True, dropout=0.5, bidirectional=False, device='cuda', 
    ...     dtype=torch.float32
    ... )
    '''
    nonlinearity: NonLinearity = NonLinearity.Tanh
    
    @overload
    def __init__(
        self, 
        nonlinearity: NonLinearity,
        input_size: int,
        hidden_size: int,
        nlays: int,
        bias: bool,
        batch_first: bool ,
        dropout: float ,
        bidirectional: bool ,
        device: deviceq,
        dtype: dtypeq,
    ): ...
    @overload
    def __init__(
        self, 
        nonlinearity: NonLinearity,
        **kwargs: BaseRecurrentKeywords
    ): ...
    def __init__(self, *args, **kwargs: 'RNNKeywords') -> 'RNNKeywords':
        '''Typed dictionary for specifying keywords specific to RNN layers, extending `BaseRecurrentKeywords`.'''
        return super().__init__(self, *(), **kwargs)

# %% ../nbs/03_kwds.ipynb 77
class RecurrentKeywords(LSTMKeywords, RNNKeywords, **_TYPE_DICT_KWDS): 
    '''Typed dictionary for specifying keywords for recurrent neural network layers, combining `LSTMKeywords` and `RNNKeywords`.

    Examples
    --------
    >>> recurrent_keywords = RecurrentKeywords(
    ...     input_size=64, hidden_size=128, nlays=2, proj_size=64, nonlinearity=NonLinearity.Tanh, 
    ...     bias=True, batch_first=True, dropout=0.5, bidirectional=False, device='cuda', 
    ...     dtype=torch.float32
    ... )
    '''
    
    @overload
    def __init__(
        self, 
        nonlinearity: NonLinearity,
        input_size: int,
        hidden_size: int,
        nlays: int,
        bias: bool,
        batch_first: bool ,
        dropout: float ,
        bidirectional: bool ,
        device: deviceq,
        dtype: dtypeq,
    ): ...
    @overload
    def __init__(
        self, 
        proj_size: int,
        input_size: int,
        hidden_size: int,
        nlays: int,
        bias: bool,
        batch_first: bool ,
        dropout: float ,
        bidirectional: bool ,
        device: deviceq,
        dtype: dtypeq,
    ): ...
    @overload
    def __init__(
        self, *args, **kwargs: LSTMKeywords | RNNKeywords
    ): ...
    def __init__(self, *args, **kwargs: 'RecurrentKeywords') -> 'RecurrentKeywords':
        '''Typed dictionary for specifying keywords for recurrent neural network layers, combining `LSTMKeywords` and `RNNKeywords`.'''
        return super().__init__(self, *(), **kwargs)

# %% ../nbs/03_kwds.ipynb 80
class RecurrentAutoEncoderKeywords(RecurrentKeywords, **_TYPE_DICT_KWDS):
    '''Typed dictionary for specifying keywords for recurrent autoencoder models, extending `RecurrentKeywords`.

    Attributes
    ----------
    output_size : int
        The size of the output layer.
        
    kind : RecurrentLayer
        The type of recurrent layer to use.

    Examples
    --------
    >>> recurrent_ae_keywords = RecurrentAutoEncoderKeywords(
    ...     input_size=64, hidden_size=128, nlays=2, proj_size=64, output_size=32, 
    ...     kind=RecurrentLayer.LSTM, nonlinearity=NonLinearity.Tanh, bias=True, batch_first=True, 
    ...     dropout=0.5, bidirectional=False, device='cuda', dtype=torch.float32
    ... )
    '''
    output_size: int
    kind: RecurrentLayer = RecurrentLayer.LSTM
    
    @overload
    def __init__(
        self, 
        output_size: int,
        kind: RecurrentLayer,
        nonlinearity: NonLinearity,
        input_size: int,
        hidden_size: int,
        nlays: int,
        bias: bool,
        batch_first: bool ,
        dropout: float ,
        bidirectional: bool ,
        device: deviceq,
        dtype: dtypeq,
    ): ...
    @overload
    def __init__(
        self, 
        output_size: int,
        kind: RecurrentLayer,
        proj_size: int,
        input_size: int,
        hidden_size: int,
        nlays: int,
        bias: bool,
        batch_first: bool ,
        dropout: float ,
        bidirectional: bool ,
        device: deviceq,
        dtype: dtypeq,
    ): ...
    @overload
    def __init__(
        self, 
        output_size: int,
        kind: RecurrentLayer,
        input_size: int,
        hidden_size: int,
        nlays: int,
        bias: bool,
        batch_first: bool ,
        dropout: float ,
        bidirectional: bool ,
        device: deviceq,
        dtype: dtypeq,
        **kwargs: LSTMKeywords | RNNKeywords
    ): ...
    @overload
    def __init__(
        self, 
        output_size: int,
        kind: RecurrentLayer,
        **kwargs: RecurrentKeywords
    ): ...
    def __init__(self, *args, **kwargs: 'RecurrentAutoEncoderKeywords') -> 'RecurrentAutoEncoderKeywords':
        '''Typed dictionary for specifying keywords specific to RNN layers, extending `BaseRecurrentKeywords`.'''
        result, params = dict(), kwargs.copy()
        kind = params.pop(KIND, RecurrentLayer.LSTM)
        if RecurrentLayer(kind) != RecurrentLayer.RNN:  params.pop(NONLINEARITY, None)
        if RecurrentLayer(kind) != RecurrentLayer.LSTM: params.pop(PROJ_SIZE, None)
        return super().__init__(self, *args, **{**result, **params})

# %% ../nbs/03_kwds.ipynb 83
class RecurrentEncoderKeywords(RecurrentAutoEncoderKeywords, **_TYPE_DICT_KWDS):
    def __init__(self, *args, **kwargs: 'RecurrentEncoderKeywords'):
        super().__init__(self, *args, **kwargs)
        input_size  = self.pop(INPUT_SIZE,  None)
        hidden_size = self.pop(HIDDEN_SIZE, None)
        output_size = self.pop(OUTPUT_SIZE, None)
        self.setdefault(INPUT_SIZE,  input_size or hidden_size)
        self.setdefault(HIDDEN_SIZE, output_size or hidden_size)
        

# %% ../nbs/03_kwds.ipynb 86
class RecurrentDecoderKeywords(RecurrentAutoEncoderKeywords, **_TYPE_DICT_KWDS):
    def __init__(self, *args, **kwargs: 'RecurrentDecoderKeywords'):
        super().__init__(self, *args, **kwargs)
        input_size  = self.pop(INPUT_SIZE,  None)
        hidden_size = self.pop(HIDDEN_SIZE, None)
        output_size = self.pop(OUTPUT_SIZE, None)
        self.setdefault(HIDDEN_SIZE, input_size  or hidden_size)
        self.setdefault(OUTPUT_SIZE, output_size or hidden_size)

# %% ../nbs/03_kwds.ipynb 89
#| export
