import math

import numpy as np

import megengine as mge
import megengine.functional as F
import megengine.module as M

# ================================= GRU Implementation ==========================================================


class GRUCell(M.Module):

    """
    An implementation of GRUCell.

    """

    def __init__(self, input_size, hidden_size, bias=True):
        super().__init__()
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.bias = bias
        self.ih = M.Linear(input_size, 3 * hidden_size, bias=bias)
        self.hh = M.Linear(hidden_size, 3 * hidden_size, bias=bias)
        self.reset_parameters()

    def reset_parameters(self):
        std = 1.0 / math.sqrt(self.hidden_size)
        for w in self.parameters():
            M.init.uniform_(w, -std, std)

    def forward(self, x, hidden):

        x = 
F.reshape(x, (-1, x.shape[1]))