How to use the chainer.Chain function in chainer

# coding: utf-8

import numpy as np
import chainer
import chainer.functions as F


class Softmax(chainer.Chain):
    def forward(self, x):
        return F.softmax(x)


class SoftmaxAxis(chainer.Chain):
    def forward(self, x):
        return F.softmax(x, axis=2)


# ======================================

from chainer_compiler.elichika import testtools
import numpy as np


def main():
    np.random.seed(314)
    a = np.random.rand(3, 5, 4).astype(np.float32)

    testtools.generate_testcase(Softmax(), [a])

import chainer.testing
import chainer.utils
import mpi4py.MPI
import numpy
import pytest

from chainermn.communicators.naive_communicator import NaiveCommunicator
from chainermn.communicators.pure_nccl_communicator import PureNcclCommunicator
from chainermn.links import MultiNodeBatchNormalization
from chainermn import nccl


mpi_comm = mpi4py.MPI.COMM_WORLD


class ModelNormalBN(chainer.Chain):
    def __init__(self, n_in=3, n_units=3, n_out=2):
        super(ModelNormalBN, self).__init__()
        with self.init_scope():
            self.l1 = chainer.links.Linear(n_in, n_units, nobias=True)
            self.bn1 = chainer.links.BatchNormalization(n_units)
            self.l2 = chainer.links.Linear(n_in, n_units, nobias=True)
            self.bn2 = chainer.links.BatchNormalization(n_units)
            self.l3 = chainer.links.Linear(n_in, n_out)
        self.train = True

    def __call__(self, x):
        h = chainer.functions.relu(self.bn1(self.l1(x)))
        h = chainer.functions.relu(self.bn2(self.l2(h)))
        return self.l3(h)

# -*- coding: utf-8 -*-

import chainer.functions as F
import chainer.links as L
import chainer
from chainer import Chain
import numpy as np
import six

class LetterClassifyer(Chain):
    def __init__(self, vocab_size, embed_size, hidden_size, class_size=2):
        super(LetterClassifyer, self).__init__(
            embed = L.EmbedID(vocab_size, embed_size),
            fc1 = L.Linear(embed_size*2, hidden_size*2),
            fc2 = L.Linear(hidden_size*2, class_size)
            )
    def forward(self, x_list):
        cl_list = []
        for x in x_list:
            wvec = self.embed(x)
            cl_list.append(wvec)
        cr_list = []
        for x in reversed(x_list):
            wvec = self.embed(x)
            cr_list.append(wvec)
        xi_list = []

Original paper: 'Inception-v4, Inception-ResNet and the Impact of Residual Connections on Learning,'
    https://arxiv.org/abs/1602.07261.
"""

__all__ = ['InceptionV4', 'inceptionv4']

import os
import chainer.functions as F
import chainer.links as L
from chainer import Chain
from functools import partial
from chainer.serializers import load_npz
from .common import SimpleSequential, Concurrent


class InceptConv(Chain):
    """
    InceptionV4 specific convolution block.

    Parameters:
    ----------
    in_channels : int
        Number of input channels.
    out_channels : int
        Number of output channels.
    ksize : int or tuple/list of 2 int
        Convolution window size.
    stride : int or tuple/list of 2 int
        Stride of the convolution.
    pad : int or tuple/list of 2 int
        Padding value for convolution layer.
    """

super(LSTMRnn, self).__init__(
            xe = SrcEmbed(vocab_size, embed_size),
            lstm = links.LSTM(embed_size, hidden_size),
            hy = links.Linear(hidden_size, vocab_size),
        )

    def reset(self):
        self.zerograds()

    def __call__(self, x):
        e = self.xe(x)
        h = self.lstm(e)
        y = self.hy(h)
        return y

class LSTMEncoder(Chain):
    def __init__(self, embed_size, hidden_size):
        super(LSTMEncoder, self).__init__(
            lstm = links.LSTM(embed_size, hidden_size),
        )
    def reset(self):
        self.zerograds()
    def __call__(self, x):
        h = self.lstm(x)
        return h

class Attention(Chain):
  def __init__(self, hidden_size, embed_size):
    super(Attention, self).__init__(
        aw = links.Linear(embed_size, hidden_size),
        pw = links.Linear(hidden_size, hidden_size),
        we = links.Linear(hidden_size, 1),

class StyleBlock(chainer.Chain):
    def __init__(self, w_in, ch):
        super().__init__()
        self.w_in = w_in
        self.ch = ch
        with self.init_scope(): 
            self.s = EqualizedLinear(w_in, ch, initial_bias=chainer.initializers.One(), gain=1)
            self.b = EqualizedLinear(w_in, ch, initial_bias=chainer.initializers.Zero(), gain=1)
    
    def __call__(self, w, h):
        ws = self.s(w)
        wb = self.b(w)
        return AdaIN(h, ws, wb)


class SynthesisBlock(chainer.Chain):
    def __init__(self, ch=512, ch_in=512, w_ch=512, upsample=True, enable_blur=False):
        super().__init__()
        self.upsample = upsample
        self.ch = ch
        self.ch_in = ch_in
        self.w_ch = w_ch
        with self.init_scope(): 
            if not upsample:
                self.W = chainer.Parameter(shape=(ch_in, 4, 4))
                self.W.data[:] = 1  # w_data_tmp

            self.b0 = L.Bias(axis=1, shape=(ch,))
            self.b1 = L.Bias(axis=1, shape=(ch,))
            self.n0 = NoiseBlock(ch)
            self.n1 = NoiseBlock(ch)

import numpy as np


class MyAdaGrad:
    def __init__(self, var, xp, eps=1e-8, lr=0.001):
        self.r = xp.ones_like(var) * eps
        self.lr = lr

    def calc_update(self, grads):
        self.r = self.r + grads * grads
        eta = F.broadcast_to(self.lr, grads.shape) / F.sqrt(self.r)
        return - eta * grads


class ResNetAuxABGenerator(chainer.Chain):
    def __init__(self, ch=64, dim_z=128, bottom_width=4, activation=F.relu, n_classes=0, distribution="normal",
                 dim_a=1000, dim_b=1000, dim_zeta=10000, T=1, learned_lr=False,
                 initial_fast_alpha=0.001, limit_fast_alpha=0.01, step_fast_alpha=0.000001):
        super(ResNetAuxABGenerator, self).__init__()
        initializer = chainer.initializers.GlorotUniform()
        self.ch = ch
        self.bottom_width = bottom_width
        self.activation = activation
        self.distribution = distribution
        self.dim_z = dim_z
        self.dim_zeta = dim_zeta
        self.n_classes = n_classes

        self.T = T
        self.learned_lr = learned_lr
        self.initial_fast_alpha = initial_fast_alpha

class SRGeneratorResBlock(chainer.Chain):
    def __init__(self):
        super().__init__(
            c1=chainer.links.Convolution2D(64, 64, ksize=3, stride=1, pad=1, wscale=0.02 * math.sqrt(64 * 3 * 3)),
            bn1=chainer.links.BatchNormalization(64),
            c2=chainer.links.Convolution2D(64, 64, ksize=3, stride=1, pad=1, wscale=0.02 * math.sqrt(64 * 3 * 3)),
            bn2=chainer.links.BatchNormalization(64),
        )

    def __call__(self, x: chainer.Variable, test=False):
        h = chainer.functions.relu(self.bn1(self.c1(x), test=test))
        h = self.bn2(self.c2(h))
        return h + x  # residual


class SRGeneratorUpScaleBlock(chainer.Chain):
    def __init__(self):
        super().__init__(
            conv=chainer.functions.Convolution2D(in_channels=64, out_channels=256, ksize=3, stride=1, pad=1,
                                                 wscale=0.02 * math.sqrt(64 * 3 * 3))
        )

    def __call__(self, x: chainer.Variable):
        h = self.conv(x)
        h = pixel_shuffle_upscale(h)
        h = chainer.functions.relu(h)
        return h


class SRGenerator(chainer.Chain):
    def __init__(self):
        super().__init__(

def fzeros(shape):
    return XP.__zeros(shape, XP.__lib.float32)

  @staticmethod
  def __array(array, dtype):
    return Variable(XP.__lib.array(array, dtype=dtype))

  @staticmethod
  def iarray(array):
    return XP.__array(array, XP.__lib.int32)

  @staticmethod
  def farray(array):
    return XP.__array(array, XP.__lib.float32)

class Encoder(Chain):
  def __init__(self, vocab_size, embed_size, hidden_size):
    super(Encoder, self).__init__(
        xe = links.EmbedID(vocab_size, embed_size),
        eh = links.Linear(embed_size, 4 * hidden_size),
        hh = links.Linear(hidden_size, 4 * hidden_size),
    )

  def __call__(self, x, c, h):
    e = functions.tanh(self.xe(x))
    return functions.lstm(c, self.eh(e) + self.hh(h)) 

class Decoder(Chain):
  def __init__(self, vocab_size, embed_size, hidden_size):
    super(Decoder, self).__init__(
        ye = links.EmbedID(vocab_size, embed_size),
        eh = links.Linear(embed_size, 4 * hidden_size),

def get_initial_states(self, mb_size):
        mb_initial_state = F.broadcast_to(F.reshape(
            self.initial_state, (1, self.out_size)), (mb_size, self.out_size))
        return (mb_initial_state,)

    def __call__(self, prev_states, x_in):
        assert len(prev_states) == 1
        prev_state = prev_states[0]
        new_state = self.gru(prev_state, x_in)
        return (new_state,)

    def get_nb_states(self):
        return 1


class FastGRUCell(Chain):
    def __init__(self, in_size, out_size, init=None, bias_init=None):
        log.info("Creating GRUCell(%i, %i)" % (in_size, out_size))
        super(FastGRUCell, self).__init__(
            gru=faster_gru.GRU(out_size, in_size, init=init, bias_init=bias_init),
        )
        self.add_param("initial_state", (1, out_size))
        self.initial_state.data[...] = self.xp.random.randn(out_size)
        self.out_size = out_size
        self.in_size = in_size

    def get_initial_states(self, mb_size):
        mb_initial_state = F.broadcast_to(F.reshape(
            self.initial_state, (1, self.out_size)), (mb_size, self.out_size))
        return (mb_initial_state,)

    def __call__(self, prev_states, x_in):