# Path: deep_architect/helpers/tensorflow_eager_support.py
# class TensorflowEagerModule(co.Module):
#     def __init__(self,
#                  name,
#                  compile_fn,
#                  name_to_hyperp,
#                  input_names,
#                  output_names,
#                  scope=None):
#     def _compile(self):
#     def _forward(self):
#     def _update(self):
# def set_is_training(outputs, is_training):
#     def fn(mx):
# def siso_tensorflow_eager_module(name, compile_fn, name_to_hyperp, scope=None):
# def siso_tensorflow_eager_module_from_tensorflow_op_fn(layer_fn,
#                                                        name_to_hyperp,
#                                                        scope=None,
#                                                        name=None):
#     def compile_fn(di, dh):
#         def forward_fn(di, is_training=False):
# def get_num_trainable_parameters():


from builtins import str
from deep_architect.helpers import tensorflow_eager_support as htfe
import tensorflow as tf
import deep_architect.modules as mo

TFEM = htfe.TensorflowEagerModule


def avg_pool(h_kernel_size, h_stride):

    def compile_fn(di, dh):

        def fn(di, is_training=True):
            with tf.device('/gpu:0'):
                return {
                    'out':
                    tf.nn.avg_pool(di['in'],
                                   [1, dh['kernel_size'], dh['kernel_size'], 1],
                                   [1, dh['stride'], dh['stride'], 1], 'SAME')
                }

        return fn


# Path: deep_architect/helpers/tensorflow_eager_support.py
# class TensorflowEagerModule(co.Module):
#     def __init__(self,
#                  name,
#                  compile_fn,
#                  name_to_hyperp,
#                  input_names,
#                  output_names,
#                  scope=None):
#     def _compile(self):
#     def _forward(self):
#     def _update(self):
# def set_is_training(outputs, is_training):
#     def fn(mx):
# def siso_tensorflow_eager_module(name, compile_fn, name_to_hyperp, scope=None):
# def siso_tensorflow_eager_module_from_tensorflow_op_fn(layer_fn,
#                                                        name_to_hyperp,
#                                                        scope=None,
#                                                        name=None):
#     def compile_fn(di, dh):
#         def forward_fn(di, is_training=False):
# def get_num_trainable_parameters():


# Path: dev/enas/search_space/common_ops.py
from builtins import str
from deep_architect.helpers import tensorflow_eager_support as htfe
import tensorflow as tf
import deep_architect.modules as mo


TFEM = htfe.TensorflowEagerModule


def avg_pool(h_kernel_size, h_stride):

    def compile_fn(di, dh):

        def fn(di, is_training=True):
            with tf.device('/gpu:0'):
                return {
                    'out':
                    tf.nn.avg_pool(di['in'],
                                   [1, dh['kernel_size'], dh['kernel_size'], 1],
                                   [1, dh['stride'], dh['stride'], 1], 'SAME')
                }

        return fn


    return htfe.siso_tensorflow_eager_module('AvgPool', compile_fn, {