#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""Parameter initialization for transducer RNN/Transformer parts."""

import six

from espnet.nets.pytorch_backend.initialization import lecun_normal_init_parameters
from espnet.nets.pytorch_backend.initialization import set_forget_bias_to_one

from espnet.nets.pytorch_backend.transformer.initializer import initialize


def initializer(model, args):
    """Initialize transducer model.

    Args:
        model (torch.nn.Module): transducer instance
        args (Namespace): argument Namespace containing options

    """
    if args.dtype != "transformer":
        if args.etype == "transformer":
            initialize(model.encoder, args.transformer_init)
            lecun_normal_init_parameters(model.dec)
        else:
            lecun_normal_init_parameters(model)

        model.dec.embed.weight.data.normal_(0, 1)

        for l in six.moves.range(len(model.dec.decoder)):
            set_forget_bias_to_one(model.dec.decoder[l].bias_ih)
    else:
        if args.etype == "transformer":
            initialize(model, args.transformer_init)
        else:
            lecun_normal_init_parameters(model.encoder)
            initialize(model.decoder, args.transformer_init)

[((24, 12, 24, 60), 'espnet.nets.pytorch_backend.transformer.initializer.initialize', 'initialize', ({(24, 23, 24, 36): 'model.encoder', (24, 38, 24, 59): 'args.transformer_init'}, {}), '(model.encoder, args.transformer_init)', False, 'from espnet.nets.pytorch_backend.transformer.initializer import initialize\n'), ((25, 12, 25, 51), 'espnet.nets.pytorch_backend.initialization.lecun_normal_init_parameters', 'lecun_normal_init_parameters', ({(25, 41, 25, 50): 'model.dec'}, {}), '(model.dec)', False, 'from espnet.nets.pytorch_backend.initialization import lecun_normal_init_parameters\n'), ((27, 12, 27, 47), 'espnet.nets.pytorch_backend.initialization.lecun_normal_init_parameters', 'lecun_normal_init_parameters', ({(27, 41, 27, 46): 'model'}, {}), '(model)', False, 'from espnet.nets.pytorch_backend.initialization import lecun_normal_init_parameters\n'), ((32, 12, 32, 64), 'espnet.nets.pytorch_backend.initialization.set_forget_bias_to_one', 'set_forget_bias_to_one', ({(32, 35, 32, 63): 'model.dec.decoder[l].bias_ih'}, {}), '(model.dec.decoder[l].bias_ih)', False, 'from espnet.nets.pytorch_backend.initialization import set_forget_bias_to_one\n'), ((35, 12, 35, 52), 'espnet.nets.pytorch_backend.transformer.initializer.initialize', 'initialize', ({(35, 23, 35, 28): 'model', (35, 30, 35, 51): 'args.transformer_init'}, {}), '(model, args.transformer_init)', False, 'from espnet.nets.pytorch_backend.transformer.initializer import initialize\n'), ((37, 12, 37, 55), 'espnet.nets.pytorch_backend.initialization.lecun_normal_init_parameters', 'lecun_normal_init_parameters', ({(37, 41, 37, 54): 'model.encoder'}, {}), '(model.encoder)', False, 'from espnet.nets.pytorch_backend.initialization import lecun_normal_init_parameters\n'), ((38, 12, 38, 60), 'espnet.nets.pytorch_backend.transformer.initializer.initialize', 'initialize', ({(38, 23, 38, 36): 'model.decoder', (38, 38, 38, 59): 'args.transformer_init'}, {}), '(model.decoder, args.transformer_init)', False, 'from espnet.nets.pytorch_backend.transformer.initializer import initialize\n')]