# Path: util/nn_cells.py
# def weight_variable(shape,
#                     init_method=None,
#                     dtype=tf.float32,
#                     init_para=None,
#                     wd=None,
#                     seed=1234,
#                     name=None,
#                     trainable=True):
#     def __init__(self,
#                  input_dim,
#                  hidden_dim,
#                  init_method,
#                  wd=None,
#                  dtype=tf.float32,
#                  init_std=None,
#                  scope="GRU"):
#     def __call__(self, x, state, is_batch_mul=False):
#     def __init__(self,
#                  input_dim, hidden_dim,
#                  init_method, wd=None,
#                  dtype=tf.float32, init_std=None,
#                  scope="SIM_GRU"):
#     def __call__(self, x, state, is_batch_mul=False):
#     def __init__(self,
#                  dims,
#                  init_method,
#                  act_func=None,
#                  add_bias=True,
#                  wd=None,
#                  dtype=tf.float32,
#                  init_std=None,
#                  scope="MLP",
#                  use_dropout=False):
#     def __call__(self, x, dropout_mask=None):
#     def __init__(self,
#                  input_dim,
#                  hidden_dim,
#                  wd=None,
#                  dtype=tf.float32,
#                  init_std=None,
#                  scope="LSTM"):
#     def __call__(self, x, state, memory):
#     def __init__(self,
#                  dims,
#                  act_func=None,
#                  add_bias=True,
#                  wd=None,
#                  dtype=tf.float32,
#                  init_std=None,
#                  scope="ResMLP"):
#     def __call__(self, x):
#     def __init__(self,
#                  message_dim,
#                  embedding_dim,
#                  hidden_shape,
#                  init_method,
#                  act_func_type,
#                  add_bias,
#                  wd=None,
#                  dtype=tf.float32,
#                  init_std=None,
#                  scope="MLPU"):
#     def __call__(self, message, hidden_embedding):
# class GRU(object):
# class simple_GRU(object):
# class MLP(object):
# class LSTM(object):
# class ResMLP(MLP):
# class MLPU(MLP):


import init_path
import tensorflow as tf
from util import nn_cells as nn

            # step 1_4: build the mlp for the propogation between nodes
            MLP_prop_shape = self._network_shape + \
                [self._hidden_dim] + [self._hidden_dim]
            self._MLP_prop = {
                i_edge: nn.MLP(
                    MLP_prop_shape,
                    init_method=self._init_method,
                    act_func=['tanh'] * (len(MLP_prop_shape) - 1),
                    add_bias=True,
                    scope='MLP_prop_edge_{}'.format(i_edge)
                )
                for i_edge in self._node_info['edge_type_list']
            }

            # step 1_5: build the node update function for each node type
            if self._node_update_method == 'GRU':
                self._Node_update = {
                    i_node_type: nn.GRU(
                        self._hidden_dim,
                        self._hidden_dim,
                        init_method=self._init_method,
                        scope='GRU_node_{}'.format(i_node_type)
                    )
                    for i_node_type in self._node_info['node_type_dict']
                }
            else:
                assert self._node_update_method == 'MLP'
                hidden_MLP_update_shape = self._network_shape
                self._Node_update = {

# Path: util/nn_cells.py
# def weight_variable(shape,
#                     init_method=None,
#                     dtype=tf.float32,
#                     init_para=None,
#                     wd=None,
#                     seed=1234,
#                     name=None,
#                     trainable=True):
#     def __init__(self,
#                  input_dim,
#                  hidden_dim,
#                  init_method,
#                  wd=None,
#                  dtype=tf.float32,
#                  init_std=None,
#                  scope="GRU"):
#     def __call__(self, x, state, is_batch_mul=False):
#     def __init__(self,
#                  input_dim, hidden_dim,
#                  init_method, wd=None,
#                  dtype=tf.float32, init_std=None,
#                  scope="SIM_GRU"):
#     def __call__(self, x, state, is_batch_mul=False):
#     def __init__(self,
#                  dims,
#                  init_method,
#                  act_func=None,
#                  add_bias=True,
#                  wd=None,
#                  dtype=tf.float32,
#                  init_std=None,
#                  scope="MLP",
#                  use_dropout=False):
#     def __call__(self, x, dropout_mask=None):
#     def __init__(self,
#                  input_dim,
#                  hidden_dim,
#                  wd=None,
#                  dtype=tf.float32,
#                  init_std=None,
#                  scope="LSTM"):
#     def __call__(self, x, state, memory):
#     def __init__(self,
#                  dims,
#                  act_func=None,
#                  add_bias=True,
#                  wd=None,
#                  dtype=tf.float32,
#                  init_std=None,
#                  scope="ResMLP"):
#     def __call__(self, x):
#     def __init__(self,
#                  message_dim,
#                  embedding_dim,
#                  hidden_shape,
#                  init_method,
#                  act_func_type,
#                  add_bias,
#                  wd=None,
#                  dtype=tf.float32,
#                  init_std=None,
#                  scope="MLPU"):
#     def __call__(self, message, hidden_embedding):
# class GRU(object):
# class simple_GRU(object):
# class MLP(object):
# class LSTM(object):
# class ResMLP(MLP):
# class MLPU(MLP):


# Path: network/pruning_network.py
import init_path
import tensorflow as tf
from util import nn_cells as nn


            # step 1_4: build the mlp for the propogation between nodes
            MLP_prop_shape = self._network_shape + \
                [self._hidden_dim] + [self._hidden_dim]
            self._MLP_prop = {
                i_edge: nn.MLP(
                    MLP_prop_shape,
                    init_method=self._init_method,
                    act_func=['tanh'] * (len(MLP_prop_shape) - 1),
                    add_bias=True,
                    scope='MLP_prop_edge_{}'.format(i_edge)
                )
                for i_edge in self._node_info['edge_type_list']
            }

            # step 1_5: build the node update function for each node type
            if self._node_update_method == 'GRU':
                self._Node_update = {
                    i_node_type: nn.GRU(
                        self._hidden_dim,
                        self._hidden_dim,
                        init_method=self._init_method,
                        scope='GRU_node_{}'.format(i_node_type)
                    )
                    for i_node_type in self._node_info['node_type_dict']
                }
            else:
                assert self._node_update_method == 'MLP'
                hidden_MLP_update_shape = self._network_shape
                self._Node_update = {

                    i_node_type: nn.MLPU(