{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e2ee9603",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow.keras as keras\n",
    "import tensorflow.keras.backend as K\n",
    "from tensorflow.keras.models import Sequential, Model\n",
    "from tensorflow.keras.layers import Input, Dense, GRU, Conv1D, Conv2D, ZeroPadding2D, Conv2DTranspose, concatenate\n",
    "from tensorflow.keras.layers import BatchNormalization, LayerNormalization\n",
    "from tensorflow.keras.activations import elu, relu, sigmoid, tanh\n",
    "from tensorflow.keras.constraints import Constraint\n",
    "# from tensorflow.nn import relu, sigmoid, tanh\n",
    "from tensorflow.keras.layers import Layer\n",
    "from keras.regularizers import l2\n",
    "from keras.initializers import HeUniform\n",
    "\n",
    "import numpy as np\n",
    "import random\n",
    "from math import gcd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "c01db5ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "import import_ipynb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "0779cea1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tensorflow.keras.models import Sequential, Model\n",
    "from tensorflow.keras.layers import Input "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "c8f048d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import os\n",
    "# os.environ[\"CUDA_DEVICE_ORDER\"] = \"PCI_BUS_ID\"   # see issue #152\n",
    "# os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"\"\n",
    "# os.environ[\"TF_KERAS\"] = \"1\"\n",
    "# # os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "1b6c06d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "reg = l2(1e-5)\n",
    "class WeightClip(Constraint):\n",
    "    '''Clips the weights incident to each hidden unit to be inside a range\n",
    "    '''\n",
    "    def __init__(self, c=2):\n",
    "        self.c = c\n",
    "\n",
    "    def __call__(self, p):\n",
    "        return K.clip(p, -self.c, self.c)\n",
    "\n",
    "    def get_config(self):\n",
    "        return {'name': self.__class__.__name__,\n",
    "            'c': self.c}\n",
    "constraint = WeightClip(0.499)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "eb77ac93",
   "metadata": {},
   "outputs": [],
   "source": [
    "def grouped_conv2d_transpose(inputs, groups, convkwargs):\n",
    "    \"\"\"Performs grouped transposed convolution.\n",
    "\n",
    "    Args:\n",
    "        inputs: A `Tensor` of shape `[batch_size, h, w, c]`.\n",
    "        filters: The number of convolutional filters.\n",
    "        kernel_size: The spatial size of the convolutional kernel.\n",
    "        strides: The convolutional stride.\n",
    "        groups: The number of groups to use in the grouped convolution step.\n",
    "            The input channel count needs to be evenly divisible by `groups`.\n",
    "    Returns:\n",
    "        A `Tensor` of shape `[batch_size, new_h, new_w, filters]`.\n",
    "    \"\"\"\n",
    "    splits = tf.split(inputs, groups, axis=-1)\n",
    "    convolved_splits = [\n",
    "        Conv2DTranspose(**convkwargs)(split) for split in splits\n",
    "    ]\n",
    "#     return concatenate(convolved_splits, -1)\n",
    "    return tf.concat(convolved_splits, -1)\n",
    "\n",
    "def grouped_conv2d(inputs, groups, convkwargs):\n",
    "    splits = tf.split(inputs, groups, axis=-1)\n",
    "    convolved_splits = [\n",
    "        Conv2D(**convkwargs)(split) for split in splits\n",
    "    ]\n",
    "#     return concatenate(convolved_splits, -1)\n",
    "    return tf.concat(convolved_splits, -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd749034",
   "metadata": {},
   "outputs": [],
   "source": [
    "def Conv2DBNAct(inputs,\n",
    "                out_ch: int,\n",
    "                k: int = 1,\n",
    "                f: int = 3,\n",
    "                fstride: int = 2,\n",
    "                lookahead: int = 0,\n",
    "                batch_norm: bool = False,\n",
    "                act = 'relu',\n",
    "                bias: bool=False,\n",
    "                padding: str=\"SAME\"):\n",
    "    if padding==\"VALID\": \n",
    "        inputs = tf.pad(inputs, tf.constant([[0,0],[2,0],[0,0],[0,0]]), \"CONSTANT\", constant_values=0)\n",
    "    conv_out = Conv2D(filters=out_ch, kernel_size=(k,f), strides=(1,fstride), use_bias=bias, padding=padding)(inputs)\n",
    "    \n",
    "    if batch_norm: conv_out = BatchNormalization()(conv_out)\n",
    "    \n",
    "    if act == 'relu': conv_out = relu(conv_out)\n",
    "    elif act == 'sigmoid': conv_out = sigmoid(conv_out)\n",
    "        \n",
    "    return conv_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e1911c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def convkxf(inputs,\n",
    "            out_ch: int,\n",
    "            k: int = 1,\n",
    "            f: int = 3,\n",
    "            fstride: int = 2,\n",
    "            lookahead: int = 0,\n",
    "            batch_norm: bool = False,\n",
    "            act = 'relu',\n",
    "            mode = \"normal\",\n",
    "            depthwise: bool = True,\n",
    "            complex_in: bool = False,\n",
    "            reshape: bool = False,\n",
    "            name: str = 'conv',\n",
    "            training=False,\n",
    "            bias: bool = False,\n",
    "            BN_type = \"normal\",\n",
    "            folding=False\n",
    "           ):\n",
    "    in_ch = inputs.get_shape()[-1]\n",
    "#     bias = batch_norm is False\n",
    "    stride = 1 if f == 1 else (1, fstride)\n",
    "    fpad = (f - 1) // 2 # freq\n",
    "        \n",
    "    if depthwise: groups = gcd(in_ch, out_ch)\n",
    "    else: groups = 1\n",
    "        \n",
    "    if in_ch % groups != 0 or out_ch % groups != 0: groups = 1   \n",
    "    \n",
    "    if mode == \"normal\":\n",
    "        convkwargs = {\n",
    "            \"filters\": out_ch, #//groups if (out_ch%groups==0) else out_ch,\n",
    "            \"groups\": groups,\n",
    "            \"kernel_size\": (k, f),\n",
    "            \"strides\": stride,\n",
    "            \"use_bias\": False,\n",
    "            \"kernel_initializer\": 'he_normal',\n",
    "#             \"kernel_regularizer\": reg,\n",
    "            \"padding\": 'valid',\n",
    "            \"name\": name\n",
    "        }\n",
    "        print(convkwargs)\n",
    "        if training:\n",
    "            paddings = tf.constant([[0,0],[k-1-lookahead,lookahead],[0,0],[0,0]])\n",
    "            pad_flag = (0, 0, k - 1 - lookahead, lookahead)\n",
    "            if any(p > 0 for p in pad_flag):\n",
    "                inputs = tf.pad(inputs, paddings, \"CONSTANT\", constant_values=0)\n",
    "            \n",
    "        if fpad>0:\n",
    "            paddings = tf.constant([[0,0],[0,0],[fpad,fpad],[0,0]])\n",
    "            inputs = tf.pad(inputs, paddings, \"CONSTANT\", constant_values=0)\n",
    "            \n",
    "        conv_out = Conv2D(**convkwargs)(inputs)\n",
    "#         if groups==1: conv_out = Conv2D(**convkwargs)(inputs)\n",
    "#         else: conv_out = grouped_conv2d(inputs, groups, convkwargs)\n",
    "\n",
    "#         if groups>1: \n",
    "#             conv_out = grouped_conv2d(inputs, groups, convkwargs)\n",
    "#         else: \n",
    "#             conv_out = Conv2D(name=name, **convkwargs)(inputs)\n",
    "        \n",
    "    elif mode == \"transposed\":\n",
    "        convkwargs = {\n",
    "            \"filters\": out_ch//groups,\n",
    "            \"kernel_size\": (k, f),\n",
    "            \"strides\": stride,\n",
    "            \"use_bias\": False,\n",
    "            \"padding\": 'same',\n",
    "            \"kernel_initializer\": 'he_normal',\n",
    "#             \"kernel_regularizer\": reg,\n",
    "        }\n",
    "        if groups>1: conv_out = grouped_conv2d_transpose(inputs, groups, convkwargs)\n",
    "        else: conv_out = Conv2DTranspose(**convkwargs)(inputs) \n",
    "\n",
    "    else:\n",
    "        raise NotImplementedError()\n",
    "        \n",
    "    if groups>1 and not (max(k,f)==1): \n",
    "        conv_out = Conv2D(out_ch, kernel_size=1, kernel_initializer='he_normal',\n",
    "                          use_bias= False, name=name + '_1x1')(conv_out)\n",
    "        \n",
    "    if batch_norm: conv_out = BatchNormalization(name = name + 'BatchNorm')(conv_out)\n",
    "    \n",
    "    if reshape: \n",
    "        shape = [tf.shape(conv_out)[l] for l in range(4)]\n",
    "        conv_out = tf.reshape(conv_out,(shape[0],shape[1],shape[2]))\n",
    "    \n",
    "    if act == 'relu': conv_out = relu(conv_out)\n",
    "    elif act == 'sigmoid': conv_out = sigmoid(conv_out)\n",
    "    else: pass\n",
    "\n",
    "    return conv_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed6043e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# def convkxf_nogroup(inputs,\n",
    "#             out_ch: int,\n",
    "#             k: int = 1,\n",
    "#             f: int = 3,\n",
    "#             fstride: int = 2,\n",
    "#             lookahead: int = 0,\n",
    "#             batch_norm: bool = False,\n",
    "#             act = 'relu',\n",
    "#             mode = \"normal\",\n",
    "#             depthwise: bool = True,\n",
    "#             complex_in: bool = False,\n",
    "#             reshape: bool = False,\n",
    "#             name: str = 'conv',\n",
    "#             training=False,\n",
    "#             bias: bool = False,\n",
    "#             BN_type = \"normal\",\n",
    "#             folding=False\n",
    "#            ):\n",
    "#     in_ch = inputs.get_shape()[-1]\n",
    "# #     bias = batch_norm is False\n",
    "#     stride = 1 if f == 1 else (1, fstride)\n",
    "#     fpad = (f - 1) // 2 # freq\n",
    "        \n",
    "#     if depthwise: groups = gcd(in_ch, out_ch)\n",
    "#     else: groups = 1\n",
    "        \n",
    "#     if in_ch % groups != 0 or out_ch % groups != 0: groups = 1  \n",
    "\n",
    "#     if mode == \"normal\":\n",
    "#         convkwargs = {\n",
    "#             \"filters\": out_ch, #//groups if (out_ch%groups==0) else out_ch,\n",
    "#             \"groups\": 1,\n",
    "#             \"kernel_size\": (k, f),\n",
    "#             \"kernel_initializer\": 'he_normal',\n",
    "#             \"strides\": stride,\n",
    "#             \"use_bias\": True if (folding and (((groups==1) and not training) or complex_in)) else bias,\n",
    "#             \"padding\": 'valid',\n",
    "#             \"name\": name\n",
    "#         }\n",
    "#         print(convkwargs)\n",
    "#         if training:\n",
    "#             paddings = tf.constant([[0,0],[k-1-lookahead,lookahead],[0,0],[0,0]])\n",
    "#             pad_flag = (0, 0, k - 1 - lookahead, lookahead)\n",
    "#             if any(p > 0 for p in pad_flag):\n",
    "#                 inputs = tf.pad(inputs, paddings, \"CONSTANT\", constant_values=0)\n",
    "            \n",
    "#         if fpad>0:\n",
    "#             paddings = tf.constant([[0,0],[0,0],[fpad,fpad],[0,0]])\n",
    "#             inputs = tf.pad(inputs, paddings, \"CONSTANT\", constant_values=0)\n",
    "            \n",
    "#         conv_out = Conv2D(**convkwargs)(inputs)\n",
    "        \n",
    "#     elif mode == \"transposed\":\n",
    "# #         groups = 1\n",
    "#         convkwargs = {\n",
    "#             \"filters\": out_ch//groups,\n",
    "#             \"kernel_size\": (k, f),\n",
    "#             \"kernel_initializer\": 'he_normal',\n",
    "#             \"strides\": stride,\n",
    "#             \"use_bias\": True if (folding and  (not groups>1 and not training) or complex_in) else bias,\n",
    "#             \"padding\": 'same',\n",
    "#         }\n",
    "#         if groups>1: \n",
    "#             conv_out = grouped_conv2d_transpose(inputs, groups, convkwargs)\n",
    "#             if folding:\n",
    "#                 conv_out = Conv2D(out_ch, kernel_size=1,\n",
    "#                               use_bias= True, name=name + '_1x1')(conv_out)\n",
    "#             else:\n",
    "#                 conv_out = Conv2D(out_ch, kernel_size=1,\n",
    "#                               use_bias= False, name=name + '_1x1')(conv_out)\n",
    "#         else:\n",
    "#             conv_out = Conv2DTranspose(**convkwargs)(inputs) \n",
    "\n",
    "#     else:\n",
    "#         raise NotImplementedError()\n",
    "        \n",
    "#     if batch_norm: \n",
    "#         if BN_type == \"normal\":\n",
    "#             conv_out = BatchNormalization(name = name + 'BatchNorm')(conv_out)\n",
    "#         if BN_type == \"range\":\n",
    "#             conv_out = RangeBN(filters=out_ch)(conv_out)\n",
    "    \n",
    "#     if reshape: \n",
    "#         shape = [tf.shape(conv_out)[l] for l in range(4)]\n",
    "#         conv_out = tf.reshape(conv_out,(shape[0],shape[1],shape[2]))\n",
    "    \n",
    "#     if act == 'relu': conv_out = relu(conv_out)\n",
    "#     elif act == 'sigmoid': conv_out = sigmoid(conv_out)\n",
    "#     else: pass\n",
    "\n",
    "#     return conv_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "403c342e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# def convkxf_custom_group(\n",
    "#             inputs,\n",
    "#             out_ch: int,\n",
    "#             k: int = 1,\n",
    "#             f: int = 3,\n",
    "#             fstride: int = 2,\n",
    "#             lookahead: int = 0,\n",
    "#             batch_norm: bool = False,\n",
    "#             act = 'relu',\n",
    "#             mode = \"normal\",\n",
    "#             depthwise: bool = True,\n",
    "#             complex_in: bool = False,\n",
    "#             reshape: bool = False,\n",
    "#             name: str = 'conv',\n",
    "#             training=False,\n",
    "#             bias: bool = False,\n",
    "#             BN_type = \"normal\",\n",
    "#             folding=False\n",
    "#            ):\n",
    "#     in_ch = inputs.get_shape()[-1]\n",
    "#     stride = 1 if f == 1 else (1, fstride)\n",
    "#     fpad = (f - 1) // 2 # freq\n",
    "        \n",
    "#     if depthwise: groups = gcd(in_ch, out_ch)\n",
    "#     else: groups = 1\n",
    "        \n",
    "#     if in_ch % groups != 0 or out_ch % groups != 0: groups = 1    \n",
    "\n",
    "#     if mode == \"normal\":\n",
    "#         convkwargs = {\n",
    "#             \"filters\": out_ch//groups if (out_ch%groups==0) else out_ch,\n",
    "# #             \"groups\": groups,\n",
    "#             \"kernel_size\": (k, f),\n",
    "#             \"strides\": stride,\n",
    "#             \"use_bias\": True if (folding and (((groups==1) and not training) or complex_in)) else bias,\n",
    "#             \"padding\": 'valid',\n",
    "# #             \"name\": name\n",
    "#         }\n",
    "#         print(convkwargs)\n",
    "#         if training:\n",
    "#             paddings = tf.constant([[0,0],[k-1-lookahead,lookahead],[0,0],[0,0]])\n",
    "#             pad_flag = (0, 0, k - 1 - lookahead, lookahead)\n",
    "#             if any(p > 0 for p in pad_flag):\n",
    "#                 inputs = tf.pad(inputs, paddings, \"CONSTANT\", constant_values=0)\n",
    "            \n",
    "#         if fpad>0:\n",
    "#             paddings = tf.constant([[0,0],[0,0],[fpad,fpad],[0,0]])\n",
    "#             inputs = tf.pad(inputs, paddings, \"CONSTANT\", constant_values=0)\n",
    "            \n",
    "#         if groups>1: conv_out = grouped_conv2d(inputs, groups, convkwargs)\n",
    "#         else: conv_out = Conv2D(**convkwargs)(inputs)\n",
    "        \n",
    "#     elif mode == \"transposed\":\n",
    "#         convkwargs = {\n",
    "#             \"filters\": out_ch//groups,\n",
    "#             \"kernel_size\": (k, f),\n",
    "#             \"strides\": stride,\n",
    "#             \"use_bias\": True if (folding and  (not groups>1 and not training) or complex_in) else bias,\n",
    "#             \"padding\": 'same',\n",
    "#         }\n",
    "#         if groups>1: \n",
    "#             conv_out = grouped_conv2d_transpose(inputs, groups, convkwargs)\n",
    "#             if folding:\n",
    "#                 conv_out = Conv2D(out_ch, kernel_size=1,\n",
    "#                               use_bias= True, name=name + '_1x1')(conv_out)\n",
    "#             else:\n",
    "#                 conv_out = Conv2D(out_ch, kernel_size=1,\n",
    "#                               use_bias= False, name=name + '_1x1')(conv_out)\n",
    "#         else: conv_out = Conv2DTranspose(**convkwargs)(inputs) \n",
    "\n",
    "#     else:\n",
    "#         raise NotImplementedError()\n",
    "        \n",
    "#     if batch_norm: \n",
    "#         if BN_type == \"normal\":\n",
    "#             conv_out = BatchNormalization(name = name + 'BatchNorm')(conv_out)\n",
    "#         if BN_type == \"range\":\n",
    "#             conv_out = RangeBN(filters=out_ch)(conv_out)\n",
    "    \n",
    "#     if reshape: \n",
    "#         shape = [tf.shape(conv_out)[l] for l in range(4)]\n",
    "#         conv_out = tf.reshape(conv_out,(shape[0],shape[1],shape[2]))\n",
    "    \n",
    "#     if act == 'relu': conv_out = relu(conv_out)\n",
    "#     elif act == 'sigmoid': conv_out = sigmoid(conv_out)\n",
    "#     else: pass\n",
    "\n",
    "#     return conv_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d942e98",
   "metadata": {},
   "outputs": [],
   "source": [
    "def GroupGRU(inputs, hidden, groups, return_sequences=True, name='GGRU', reshape=True, count=0, training=True, norm=False):\n",
    "    \n",
    "    grukwargs = {\n",
    "        \"units\": hidden//groups,\n",
    "        \"return_sequences\": True,\n",
    "#         \"kernel_initializer\": 'he_normal',\n",
    "#         \"kernel_constraint\": constraint, \n",
    "#         \"recurrent_constraint\": constraint, \n",
    "#         \"bias_constraint\": constraint,\n",
    "#         \"kernel_regularizer\": reg,\n",
    "#         \"recurrent_regularizer\": reg\n",
    "    }\n",
    "    if groups>1:\n",
    "        gru_list = []\n",
    "        groups_inputs = tf.split(inputs, groups, axis=-1)\n",
    "        for group_inputs in groups_inputs: # group number\n",
    "            if not training: \n",
    "                gru_tmp = GRU(**grukwargs, unroll=True, name=name + '_' +str(count))(group_inputs)\n",
    "                if norm: gru_tmp = BatchNormalization()(gru_tmp)\n",
    "            else: \n",
    "                gru_tmp = GRU(**grukwargs, name=name + '_' +str(count))(group_inputs)\n",
    "                if norm: gru_tmp = BatchNormalization()(gru_tmp)\n",
    "            gru_list.append(tf.expand_dims(gru_tmp,-1))\n",
    "            count+=1\n",
    "            \n",
    "        #gru_output = tf.stack(gru_list, axis=-1, name=name + '_stack')\n",
    "        gru_output = concatenate(gru_list, axis=-1, name=name + '_stack')\n",
    "#         gru_output = tf.concat(gru_list, axis=-1, name=name + '_stack')\n",
    "        if reshape:\n",
    "            gru_output = tf.transpose(gru_output, [0, 1, 3, 2], name=name + '_transpose')\n",
    "\n",
    "        shape = [tf.shape(gru_output)[l] for l in range(4)]\n",
    "        gru_layer_output = tf.reshape(gru_output, shape = [shape[0], shape[1], shape[-1]*shape[-2]], name=name + '_reshape')\n",
    "    else:\n",
    "        if not training: \n",
    "            gru_layer_output = GRU(**grukwargs, unroll=True, name=name)(inputs)\n",
    "            if norm: gru_layer_output = BatchNormalization()(gru_layer_output)\n",
    "        else: \n",
    "            gru_layer_output = GRU(**grukwargs, name=name)(inputs)\n",
    "            if norm: gru_layer_output = BatchNormalization()(gru_layer_output)\n",
    "    \n",
    "    gru_layer_output = relu(gru_layer_output)\n",
    "    \n",
    "    if return_sequences:\n",
    "        return gru_layer_output\n",
    "    else:\n",
    "        return gru_layer_output[:,-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10443262",
   "metadata": {},
   "outputs": [],
   "source": [
    "def GroupGRULayer(inputs, hidden, groups, num_layer, name='GGRU', add_output=True, training=True, norm=False):\n",
    "    for i in range(num_layer):\n",
    "        if i < num_layer-1: reshape = True\n",
    "        else: reshape = False\n",
    "        \n",
    "        inputs = GroupGRU(inputs, hidden, groups, return_sequences=True, name=name + str(i), \\\n",
    "                          reshape=reshape, training=training, norm=norm)\n",
    "        if add_output:\n",
    "            if i == 0 : output = inputs\n",
    "            else: output += inputs\n",
    "                \n",
    "        else: output = inputs\n",
    "        \n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8dd54c86",
   "metadata": {},
   "outputs": [],
   "source": [
    "def GroupGRU_stateful(inputs, hidden, groups, return_sequences=True, name='GGRU', reshape=True, count=0, training=True, norm=False):\n",
    "    \n",
    "    grukwargs = {\n",
    "        \"units\": hidden//groups,\n",
    "        \"return_sequences\": True,\n",
    "        \"stateful\": True,\n",
    "#         \"kernel_initializer\": 'he_normal',\n",
    "#         \"kernel_constraint\": constraint, \n",
    "#         \"recurrent_constraint\": constraint, \n",
    "#         \"bias_constraint\": constraint,\n",
    "#         \"kernel_regularizer\": reg,\n",
    "#         \"recurrent_regularizer\": reg\n",
    "    }\n",
    "    if groups>1:\n",
    "        gru_list = []\n",
    "        groups_inputs = tf.split(inputs, groups, axis=-1)\n",
    "        for group_inputs in groups_inputs: # group number\n",
    "            if not training: \n",
    "                gru_tmp = GRU(**grukwargs, unroll=True, name=name + '_' +str(count))(group_inputs)\n",
    "                if norm: gru_tmp = BatchNormalization()(gru_tmp)\n",
    "            else: \n",
    "                gru_tmp = GRU(**grukwargs, name=name + '_' +str(count))(group_inputs)\n",
    "                if norm: gru_tmp = BatchNormalization()(gru_tmp)\n",
    "            gru_list.append(tf.expand_dims(gru_tmp,-1))\n",
    "            count+=1\n",
    "            \n",
    "        gru_output = concatenate(gru_list, axis=-1, name=name + '_stack')\n",
    "        if reshape:\n",
    "            gru_output = tf.transpose(gru_output, [0, 1, 3, 2], name=name + '_transpose')\n",
    "\n",
    "        shape = [tf.shape(gru_output)[l] for l in range(4)]\n",
    "        gru_layer_output = tf.reshape(gru_output, shape = [shape[0], shape[1], shape[-1]*shape[-2]], name=name + '_reshape')\n",
    "    else:\n",
    "        gru_layer_output = GRU(**grukwargs, name=name)(inputs)\n",
    "        if norm: gru_layer_output = BatchNormalization()(gru_layer_output)\n",
    "    \n",
    "    gru_layer_output = relu(gru_layer_output)\n",
    "    \n",
    "    if return_sequences:\n",
    "        return gru_layer_output\n",
    "    else:\n",
    "        return gru_layer_output[:,-1]\n",
    "    \n",
    "def GroupGRULayer_stateful(inputs, hidden, groups, num_layer, name='GGRU', add_output=True, training=True, norm=False):\n",
    "    for i in range(num_layer):\n",
    "        if i < num_layer-1: reshape = True\n",
    "        else: reshape = False\n",
    "        \n",
    "        inputs = GroupGRU_stateful(inputs, hidden, groups, return_sequences=True, name=name + str(i), \\\n",
    "                          reshape=reshape, training=training, norm=norm)\n",
    "        if add_output:\n",
    "            if i == 0 : output = inputs\n",
    "            else: output += inputs\n",
    "                \n",
    "        else: output = inputs\n",
    "        \n",
    "    return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2af3ab74",
   "metadata": {},
   "outputs": [],
   "source": [
    "def GroupFC(inputs, hidden, groups=8, activation=None, name='GFC', shuffle=True, count=0, norm=False, training=True):\n",
    "    if not training:\n",
    "        shape = [tf.shape(inputs)[l] for l in range(3)]\n",
    "        inputs = tf.reshape(inputs,(shape[1],shape[2]))\n",
    "    \n",
    "    FC_list = []\n",
    "    fckwargs = {\n",
    "    \"units\": hidden//groups,\n",
    "#     \"kernel_initializer\": 'he_normal',\n",
    "#     \"kernel_regularizer\": reg,\n",
    "#     \"kernel_constraint\": constraint, \n",
    "#     \"bias_constraint\": constraint,\n",
    "#     \"use_bias\": False,\n",
    "    }\n",
    "    if groups>1:\n",
    "        groups_inputs = tf.split(inputs, groups, axis=-1)\n",
    "        for group_inputs in groups_inputs: # group number\n",
    "            FC_tmp = Dense(**fckwargs, name=name + '_' +str(count))(group_inputs)\n",
    "            if training: FC_list.append(FC_tmp)\n",
    "            else: FC_list.append(tf.expand_dims(FC_tmp,-1))\n",
    "            count+=1\n",
    "\n",
    "        if training: FC_output = tf.stack(FC_list, axis=-1, name=name + '_stack')\n",
    "        else: \n",
    "            FC_output = concatenate(FC_list, axis=-1)\n",
    "            FC_output = tf.expand_dims(FC_output,0)\n",
    "        \n",
    "        if shuffle:\n",
    "            FC_output = tf.transpose(FC_output, [0, 1, 3, 2], name=name + '_transpose')\n",
    "\n",
    "        shape = [tf.shape(FC_output)[l] for l in range(4)]\n",
    "        FC_layer_output = tf.reshape(FC_output, shape = [shape[0], shape[1], shape[-1]*shape[-2]], name=name + '_reshape')\n",
    "        if norm: FC_layer_output = BatchNormalization(name = name + '_' +str(count) + 'BatchNorm')(FC_layer_output)\n",
    "    else:\n",
    "        FC_layer_output = Dense(**fckwargs, name=name)(inputs)\n",
    "        \n",
    "    if activation == 'relu': FC_layer_output = relu(FC_layer_output)\n",
    "    elif activation == 'sigmoid': FC_layer_output = sigmoid(FC_layer_output)\n",
    "    elif activation == 'tanh': FC_layer_output = tanh(FC_layer_output)\n",
    "    else: FC_layer_output = FC_layer_output\n",
    "        \n",
    "    return FC_layer_output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5758b06",
   "metadata": {},
   "outputs": [],
   "source": [
    "def GroupFC_conv1d(inputs, hidden, groups=8, activation=None):\n",
    "    conv_out=Conv1D(filters=hidden,\n",
    "                    kernel_size=1,\n",
    "                    padding=\"same\",\n",
    "                    groups=groups)(inputs)\n",
    "    conv_out=Conv1D(filters=hidden,\n",
    "                    kernel_size=1,\n",
    "                    padding=\"same\",\n",
    "                    groups=1,\n",
    "                    use_bias= False)(conv_out)\n",
    "    conv_out = BatchNormalization()(conv_out)\n",
    "    \n",
    "    if activation == 'relu': conv_out = relu(conv_out)\n",
    "    elif activation == 'sigmoid': conv_out = sigmoid(conv_out)\n",
    "    else: pass\n",
    "    \n",
    "    return conv_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47b563ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "def GroupGRU_lite(inputs, hidden, groups, state=None, return_sequences=True, name='GGRU', reshape=False, \n",
    "                  count=0, training=True, norm=False):\n",
    "    grukwargs = {\n",
    "        \"units\": hidden//groups,\n",
    "        \"return_sequences\": True,\n",
    "        \"kernel_initializer\": 'he_normal',\n",
    "    }\n",
    "    if groups == 1:\n",
    "        gru_output, state_output = GRU(**grukwargs, return_state=True, unroll=True,\n",
    "                                     name=name)([inputs, state])\n",
    "\n",
    "    else:\n",
    "        groups_inputs = tf.split(inputs, groups, axis=-1)\n",
    "        gru_list, state_list = [], []\n",
    "        for group_inputs in groups_inputs: # group number\n",
    "            gru_tmp, gru_state = GRU(**grukwargs, return_state=True, unroll=True,\n",
    "                                         name=name + '_' +str(count))([group_inputs, state[count]])\n",
    "                \n",
    "        gru_list.append(gru_tmp)\n",
    "        state_list.append(gru_state)\n",
    "        \n",
    "        count+=1\n",
    "        \n",
    "        gru_output = tf.stack(gru_list, axis=-1, name=name + '_stack')\n",
    "        state_output = concatenate(state_list, axis=0)\n",
    "#         state_output = tf.concat(state_list, axis=0)\n",
    "    \n",
    "        if reshape:\n",
    "            gru_output = tf.transpose(gru_output, [0, 1, 3, 2], name=name + '_transpose')\n",
    "\n",
    "        shape = [tf.shape(gru_output)[l] for l in range(4)]\n",
    "        gru_output = tf.reshape(gru_output, shape = [shape[0], shape[1], shape[-1]*shape[-2]], name=name + '_reshape')\n",
    "        \n",
    "    if norm: gru_output = BatchNormalization()(gru_output)\n",
    "    gru_output = relu(gru_output)\n",
    "    \n",
    "    if return_sequences:\n",
    "        return gru_output, state_output\n",
    "    else:\n",
    "        return gru_output[:,-1], state_output\n",
    "    \n",
    "def GroupGRULayer_lite(inputs, hidden, groups, num_layer, state=None, name='GGRU', add_output=True, training=True, norm=False):\n",
    "    state_list = []\n",
    "    for i in range(num_layer):\n",
    "        if i < num_layer-1: reshape = True\n",
    "        else: reshape = False\n",
    "        \n",
    "        inputs, state1 = GroupGRU_lite(inputs, hidden, groups, state=state[i], \n",
    "                              return_sequences=True, name=name + str(i), \n",
    "                              reshape=reshape, training=training, norm=norm)\n",
    "        \n",
    "        state_list.append(state1) \n",
    "        \n",
    "        if add_output:\n",
    "            if i == 0 : output = inputs\n",
    "            else: output += inputs\n",
    "                \n",
    "        else: output = inputs\n",
    "        \n",
    "    return output, state_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d20a29e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d49542d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc813bc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "class RangeBN(tf.keras.layers.Layer):\n",
    "    def __init__(self,\n",
    "                 filters,\n",
    "                 momentum=0.1,\n",
    "                 num_chunks=16,\n",
    "                 eps=1e-6,\n",
    "                 name=None,\n",
    "                 **kwargs):\n",
    "        super(RangeBN, self).__init__(**kwargs)\n",
    "        \"\"\"\n",
    "            Follow the reference https://arxiv.org/abs/1805.11046, Scalable Methods for 8-bit Training of Neural Networks\n",
    "            With the range of min(x) - max(x) of  input distribution, making more tolerant to quantization.    \n",
    "            Formula: x = (x - mu)/(C(n) - range(x - mu))\n",
    "        \"\"\"\n",
    "        self.filters = filters\n",
    "        self.momentum = momentum\n",
    "        self.num_chunks = num_chunks\n",
    "        self.eps = eps\n",
    "\n",
    "    def build(self, input_shape):\n",
    "        initializer = tf.keras.initializers.RandomUniform(minval=0., maxval=1)\n",
    "        input_shape = tf.TensorShape(input_shape)\n",
    "        self.weight = tf.Variable(initial_value=initializer(\n",
    "            shape=(input_shape[-1], ), dtype=tf.dtypes.float32),\n",
    "                                  name='weights',\n",
    "                                  trainable=True)\n",
    "        self.bias = self.add_weight(\n",
    "            shape=(input_shape[-1], ),\n",
    "            initializer=tf.keras.initializers.Constant(0.),\n",
    "            name='bias',\n",
    "            trainable=True)\n",
    "        self.moving_mean = tf.Variable(initial_value=tf.constant(\n",
    "            0., shape=(self.filters, )),\n",
    "                                       shape=self.filters,\n",
    "                                       dtype=tf.dtypes.float32,\n",
    "                                       name='moving_mean',\n",
    "                                       trainable=False)\n",
    "\n",
    "        self.moving_variance = tf.Variable(initial_value=tf.constant(\n",
    "            0., shape=(self.filters, )),\n",
    "                                           shape=(self.filters, ),\n",
    "                                           dtype=tf.dtypes.float32,\n",
    "                                           name='moving_variance',\n",
    "                                           trainable=False)\n",
    "        super(RangeBN, self).build(input_shape)\n",
    "\n",
    "    def call(self, inputs, training=False):\n",
    "        if training:\n",
    "            B, H, W, C = [tf.shape(inputs)[i] for i in range(4)]\n",
    "            y = tf.transpose(inputs, (3, 0, 1, 2))\n",
    "            y = tf.reshape(y,\n",
    "                           [C, self.num_chunks, B * H * W // self.num_chunks])\n",
    "            mean_max = tf.math.reduce_max(y, [-1])\n",
    "            mean_max = tf.math.reduce_mean(mean_max, axis=-1)  # C\n",
    "            mean_min = tf.math.reduce_min(y, [-1])\n",
    "            mean_min = tf.math.reduce_mean(mean_min, axis=-1)  # C\n",
    "\n",
    "            mean = tf.math.reduce_mean(tf.reshape(y, (C, -1)), axis=-1)  # C\n",
    "            B, H, W, C = inputs.get_shape().as_list()\n",
    "            if B is None:\n",
    "                B = 1\n",
    "            upper = (0.5 * 0.35) * (1. + (math.pi * math.log(4))**0.5)\n",
    "            lower = ((2. * math.log(B * H * W // self.num_chunks))**0.5)\n",
    "            scale_fix = upper / lower\n",
    "            scale = 1 / ((mean_max - mean_min) * scale_fix + self.eps)\n",
    "\n",
    "            self.moving_mean.assign(self.moving_mean * self.momentum +\n",
    "                                    (mean * (1 - self.momentum)))\n",
    "            self.moving_variance.assign(self.moving_variance * self.momentum +\n",
    "                                        (scale * (1 - self.momentum)))\n",
    "        else:\n",
    "            mean = self.moving_mean\n",
    "            scale = self.moving_variance\n",
    "        out = (inputs - tf.reshape(mean, (1, 1, 1, self.filters))) * \\\n",
    "            tf.reshape(scale, (1, 1, 1,self.filters))\n",
    "        if self.weight is not None:\n",
    "            out = out * tf.reshape(self.weight, (1, 1, 1, self.filters))\n",
    "        if self.bias is not None:\n",
    "            out = out + tf.reshape(self.bias, (1, 1, 1, self.filters))\n",
    "        return out"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
