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

# %% auto 0
__all__ = ['recurrent_kwds', 'recurrent_encoder_kwds', 'recurrent_decoder_kwds']

# %% ../nbs/20_kwds.ipynb 6
#| export


# %% ../nbs/20_kwds.ipynb 8
#| export

# %% ../nbs/20_kwds.ipynb 11
#| export

# %% ../nbs/20_kwds.ipynb 13
#| export


# %% ../nbs/20_kwds.ipynb 15
try:
    import torch, torch.nn as nn
except ImportError:
    ...

# %% ../nbs/20_kwds.ipynb 17
#| export


# %% ../nbs/20_kwds.ipynb 19
from atyp import Tensor, DeviceQ, DTypeQ

# %% ../nbs/20_kwds.ipynb 21
from .enum import Channels, NonLinearity, RecurrentLayer

# %% ../nbs/20_kwds.ipynb 23
def recurrent_kwds(
    input_size: int, hidden_size: int, num_layers: int = 1, 
    bias: bool = True, batch_first: bool = True, dropout: float = 0.2, 
    bidirectional: bool = False, proj_size: int = 0, 
    nonlinearity: NonLinearity = NonLinearity.Tanh, 
    device: DeviceQ = None,  dtype: DTypeQ = None, 
    kind: RecurrentLayer = RecurrentLayer.LSTM,
):
    params = dict(
        input_size=input_size, hidden_size=hidden_size, num_layers=num_layers, bias=bias, 
        batch_first=batch_first, dropout=dropout, bidirectional=bidirectional,
        proj_size=proj_size, nonlinearity=nonlinearity, device=device, dtype=dtype,
    )
    
    if RecurrentLayer(kind) != RecurrentLayer.RNN:  params.pop('nonlinearity', None)
    if RecurrentLayer(kind) != RecurrentLayer.LSTM: params.pop('proj_size', None)

    return params

# %% ../nbs/20_kwds.ipynb 26
def recurrent_encoder_kwds(**kwargs):
    result, params = dict(), kwargs.copy()
    input_size  = params.pop('input_size',  None)
    hidden_size = params.pop('hidden_size', None)
    output_size = params.pop('output_size', None)
    result.setdefault('input_size',  input_size or hidden_size)
    result.setdefault('hidden_size', output_size or hidden_size)
    return {**result, **params}

# %% ../nbs/20_kwds.ipynb 28
def recurrent_decoder_kwds(**kwargs):
    result, params = dict(), kwargs.copy()
    input_size  = params.pop('input_size',  None)
    hidden_size = params.pop('hidden_size', None)
    output_size = params.pop('output_size', None)
    result.setdefault('hidden_size', input_size  or hidden_size)
    result.setdefault('output_size', output_size or hidden_size)
    return {**result, **params}
