#!/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)

[((661, 709), 'espnet.nets.pytorch_backend.transformer.initializer.initialize', 'initialize', (['model.encoder', 'args.transformer_init'], {}), '(model.encoder, args.transformer_init)\n', (671, 709), False, 'from espnet.nets.pytorch_backend.transformer.initializer import initialize\n'), ((722, 761), 'espnet.nets.pytorch_backend.initialization.lecun_normal_init_parameters', 'lecun_normal_init_parameters', (['model.dec'], {}), '(model.dec)\n', (750, 761), False, 'from espnet.nets.pytorch_backend.initialization import lecun_normal_init_parameters\n'), ((788, 823), 'espnet.nets.pytorch_backend.initialization.lecun_normal_init_parameters', 'lecun_normal_init_parameters', (['model'], {}), '(model)\n', (816, 823), False, 'from espnet.nets.pytorch_backend.initialization import lecun_normal_init_parameters\n'), ((946, 998), 'espnet.nets.pytorch_backend.initialization.set_forget_bias_to_one', 'set_forget_bias_to_one', (['model.dec.decoder[l].bias_ih'], {}), '(model.dec.decoder[l].bias_ih)\n', (968, 998), False, 'from espnet.nets.pytorch_backend.initialization import set_forget_bias_to_one\n'), ((1061, 1101), 'espnet.nets.pytorch_backend.transformer.initializer.initialize', 'initialize', (['model', 'args.transformer_init'], {}), '(model, args.transformer_init)\n', (1071, 1101), False, 'from espnet.nets.pytorch_backend.transformer.initializer import initialize\n'), ((1128, 1171), 'espnet.nets.pytorch_backend.initialization.lecun_normal_init_parameters', 'lecun_normal_init_parameters', (['model.encoder'], {}), '(model.encoder)\n', (1156, 1171), False, 'from espnet.nets.pytorch_backend.initialization import lecun_normal_init_parameters\n'), ((1184, 1232), 'espnet.nets.pytorch_backend.transformer.initializer.initialize', 'initialize', (['model.decoder', 'args.transformer_init'], {}), '(model.decoder, args.transformer_init)\n', (1194, 1232), False, 'from espnet.nets.pytorch_backend.transformer.initializer import initialize\n')]