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(): D = Discrete def conv2D(filter_size, name, weight_sharer, out_filters=None):

    def compile_fn(di, dh):
        (_, _, _, channels) = di['in'].get_shape().as_list()
        channels = channels if out_filters is None else out_filters

        conv_fn = lambda: tf.keras.layers.Conv2D(
            channels, filter_size, padding='same')
        conv = weight_sharer.get(name + '_conv_' + str(filter_size),
                                 conv_fn, lambda layer: layer.get_weights())
        if not conv.built:
            with tf.device('/gpu:0'):
                conv.build(di['in'].get_shape())
                weights = weight_sharer.load_weights(name)
                if weights is not None:
                    conv.set_weights(weights)

        def fn(di, is_training=True):
            with tf.device('/gpu:0'):
                return {'out': conv(di['in'])}

        return fn

    return htfe.siso_tensorflow_eager_module('Conv2D', compile_fn, {}) def conv2D_depth_separable(filter_size, name, weight_sharer, out_filters=None):

    def compile_fn(di, dh):
        (_, _, _, channels) = di['in'].get_shape().as_list()
        channels = channels if out_filters is None else out_filters
        conv_fn = lambda: tf.keras.layers.SeparableConv2D(
            channels, filter_size, padding='same')
        conv = weight_sharer.get(name + '_dsep_' + str(filter_size),
                                 conv_fn, lambda layer: layer.get_weights())
        if not conv.built:
            with tf.device('/gpu:0'):
                conv.build(di['in'].get_shape())
                weights = weight_sharer.load_weights(name)
                if weights is not None:
                    conv.set_weights(weights)

        def fn(di, is_training=True):
            with tf.device('/gpu:0'):
                return {'out': conv(di['in'])}

        return fn

    return htfe.siso_tensorflow_eager_module('Conv2DSeparable', compile_fn, {}) def global_pool():

    def compile_fn(di, dh):

        def fn(di, is_training):
            with tf.device('/gpu:0'):
                return {'out': tf.reduce_mean(di['in'], [1, 2])}

        return fn

    return htfe.siso_tensorflow_eager_module('GlobalPool', compile_fn, {}) def dropout(keep_prob):

    def compile_fn(di, dh):

        def fn(di, is_training=True):
            if is_training:
                with tf.device('/gpu:0'):
                    out = tf.nn.dropout(di['in'], keep_prob)
            else:
                out = di['in']
            return {'out': out}

        return fn

    return htfe.siso_tensorflow_eager_module('Dropout', compile_fn, {}) def fc_layer(num_classes, name, weight_sharer):
    name = name + '_fc_layer_' + str(num_classes)

    def compile_fn(di, dh):
        fc = weight_sharer.get(name, lambda: tf.keras.layers.Dense(num_classes),
                               lambda layer: layer.get_weights())
        if not fc.built:
            with tf.device('/gpu:0'):
                fc.build(di['in'].get_shape())
                weights = weight_sharer.load_weights(name)
                if weights is not None:
                    fc.set_weights(weights)

        def fn(di, is_training=True):
            with tf.device('/gpu:0'):
                return {'out': fc(di['in'])}

        return fn

    return htfe.siso_tensorflow_eager_module('FC_Layer', compile_fn, {}) def wrap_batch_norm_relu(io_pair,
                         add_relu=True,
                         add_bn=True,
                         weight_sharer=None,
                         name=None):
    assert add_relu or add_bn
    elements = [True, add_bn, add_relu]
    module_fns = [
        lambda: io_pair, lambda: keras_batch_normalization(
            name=name, weight_sharer=weight_sharer), relu
    ]
    return mo.siso_sequential(
        [module_fn() for i, module_fn in enumerate(module_fns) if elements[i]]) 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, {
        'kernel_size': h_kernel_size,
        'stride': h_stride,
    }) def max_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.max_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('MaxPool2D', compile_fn, {
        'kernel_size': h_kernel_size,
        'stride': h_stride,
    }) def keras_batch_normalization(name='default', weight_sharer=None):
    name = name + '_bn'

    def compile_fn(di, dh):
        bn = weight_sharer.get(name, tf.keras.layers.BatchNormalization, lambda
                               layer: layer.get_weights())
        if not bn.built:
            with tf.device('/gpu:0'):
                bn.build(di['in'].get_shape())
                weights = weight_sharer.load_weights(name)
                if weights is not None:
                    bn.set_weights(weights)

        def fn(di, is_training):
            with tf.device('/gpu:0'):
                return {'out': bn(di['in'], training=is_training)}

        return fn

    return htfe.siso_tensorflow_eager_module('BatchNormalization', compile_fn,
                                             {})
from collections import OrderedDict
from deep_architect.helpers import tensorflow_eager_support as htfe
from deep_architect.hyperparameters import D
from dev.enas.search_space.common_ops import (conv2D, conv2D_depth_separable,
                                              global_pool, dropout, fc_layer,
                                              wrap_batch_norm_relu, avg_pool,
                                              max_pool,
                                              keras_batch_normalization)
import tensorflow as tf
import numpy as np
import deep_architect.modules as mo
"""
Search space from Efficient Neural Architecture Search (Pham'17)
"""
from __future__ import print_function



TFEM = htfe.TensorflowEagerModule


class WeightSharer:

    def __init__(self, isSharing):
        self.name_to_weight = {}
        self.name_to_np_fn = {}
        self.weight_dict = {}
        self.isSharing = isSharing

    def get(self, name, construct_fn, np_fn):
        if self.isSharing:
            if name not in self.name_to_weight:
                with tf.device('/gpu:0'):
                    self.name_to_weight[name] = construct_fn()
                    self.name_to_np_fn[name] = np_fn
                print(name)

    #        self.weights_used.add(name)
    #        self.name_to_weight[name].gpu()
            return self.name_to_weight[name]
        return construct_fn()

    def load_weights(self, name):
        if name in self.weight_dict:
            return self.weight_dict[name]
        else:
            return None

    def save(self, filename):
        weight_dict = self.weight_dict
        for name in self.name_to_weight:
            weight_dict[name] = self.name_to_np_fn[name](
                self.name_to_weight[name])
        np.save(filename, weight_dict)

    def load(self, filename):
        self.weight_dict = np.load(filename).item()


# Take in array of boolean hyperparams, concatenate layers corresponding to true
# to form skip connections
def concatenate_skip_layers(h_connects, weight_sharer):

    def compile_fn(di, dh):

        def fn(di, is_training=True):
            inputs = [
                di['in' + str(i)]
                for i in range(len(dh))
                if dh['select_' + str(i)]
            ]
            inputs.append(di['in' + str(len(dh))])
            with tf.device('/gpu:0'):
                out = tf.add_n(inputs)
                return {'out': tf.add_n(inputs)}

        return fn

    return TFEM(
        'SkipConcat',
        {'select_' + str(i): h_connects[i] for i in range(len(h_connects))},
        compile_fn, ['in' + str(i) for i in range(len(h_connects) + 1)],
        ['out']).get_io()


def enas_conv(out_filters, filter_size, separable, weight_sharer, name):
    io_pair = (conv2D_depth_separable(filter_size, name, weight_sharer)

               if separable else conv2D(filter_size, name, weight_sharer))