def relu():

    def compile_fn(di, dh):
        relu = tf.keras.layers.ReLU()

        def forward_fn(di, is_training=True):
            return {'out': relu(di['in'])}

        return forward_fn

    return siso_tensorflow_eager_module('ReLU', compile_fn, {}) def batch_normalization():

    def compile_fn(di, dh):
        bn = tf.keras.layers.BatchNormalization(momentum=.9, epsilon=1e-5)

        def forward_fn(di, is_training):
            return {'out': bn(di['in'])}

        return forward_fn

    return siso_tensorflow_eager_module('BatchNormalization', compile_fn, {}) def conv2d(h_num_filters,
           h_filter_width,
           h_stride=1,
           h_dilation_rate=1,
           h_use_bias=True,
           h_padding='SAME'):

    def compile_fn(di, dh):
        conv = tf.keras.layers.Conv2D(dh['num_filters'],
                                      dh['filter_width'],
                                      dh['stride'],
                                      use_bias=dh['use_bias'],
                                      dilation_rate=dh['dilation_rate'],
                                      padding=dh['padding'])

        def forward_fn(di, is_training=True):
            return {'out': conv(di['in'])}

        return forward_fn

    return siso_tensorflow_eager_module(
        'Conv2D', compile_fn, {
            'num_filters': h_num_filters,
            'filter_width': h_filter_width,
            'stride': h_stride,
            'dilation_rate': h_dilation_rate,
            'use_bias': h_use_bias,
            'padding': h_padding
        }) def separable_conv2d(h_num_filters,
                     h_filter_width,
                     h_stride=1,
                     h_dilation_rate=1,
                     h_depth_multiplier=1,
                     h_use_bias=True,
                     h_padding='SAME'):

    def compile_fn(di, dh):

        conv_op = tf.keras.layers.SeparableConv2D(
            dh['num_filters'],
            dh['filter_width'],
            strides=dh['stride'],
            dilation_rate=dh['dilation_rate'],
            depth_multiplier=dh['depth_multiplier'],
            use_bias=dh['use_bias'],
            padding=dh['padding'])

        def fn(di, is_training=True):
            return {'out': conv_op(di['in'])}

        return fn

    return siso_tensorflow_eager_module(
        'SeparableConv2D', compile_fn, {
            'num_filters': h_num_filters,
            'filter_width': h_filter_width,
            'stride': h_stride,
            'use_bias': h_use_bias,
            'dilation_rate': h_dilation_rate,
            'depth_multiplier': h_depth_multiplier,
            'padding': h_padding
        }) def avg_pool2d(h_kernel_size, h_stride=1, h_padding='SAME'):

    def compile_fn(di, dh):
        pool = tf.keras.layers.AveragePooling2D(dh['kernel_size'],
                                                dh['stride'],
                                                padding=dh['padding'])

        def forward_fn(di, is_training=True):

            return {'out': pool(di['in'])}

        return forward_fn

    return siso_tensorflow_eager_module('MaxPool2D', compile_fn, {
        'kernel_size': h_kernel_size,
        'stride': h_stride,
        'padding': h_padding
    }) def max_pool2d(h_kernel_size, h_stride=1, h_padding='SAME'):

    def compile_fn(di, dh):
        pool = tf.keras.layers.MaxPooling2D(dh['kernel_size'],
                                            dh['stride'],
                                            padding=dh['padding'])

        def forward_fn(di, is_training=True):
            return {'out': pool(di['in'])}

        return forward_fn

    return siso_tensorflow_eager_module('MaxPool2D', compile_fn, {
        'kernel_size': h_kernel_size,
        'stride': h_stride,
        'padding': h_padding
    }) def min_pool2d(h_kernel_size, h_stride=1, h_padding='SAME'):

    def compile_fn(di, dh):
        pool = tf.keras.layers.MaxPooling2D(dh['kernel_size'],
                                            dh['stride'],
                                            padding=dh['padding'])
        negate = tf.keras.layers.Lambda(lambda x: -1 * x, name='negate')

        def forward_fn(di, is_training=True):
            return {'out': negate(pool(negate(di['in'])))}

        return forward_fn

    return siso_tensorflow_eager_module('MinPool2D', compile_fn, {
        'kernel_size': h_kernel_size,
        'stride': h_stride,
        'padding': h_padding
    }) def fc_layer(h_num_units):

    def compile_fn(di, dh):
        fc = tf.keras.layers.Dense(dh['num_units'])

        def forward_fn(di, is_training=True):
            return {'out': fc(di['in'])}

        return forward_fn

    return siso_tensorflow_eager_module('FCLayer', compile_fn,
                                        {'num_units': h_num_units}) def global_pool2d():

    def compile_fn(di, dh):
        pool = tf.keras.layers.GlobalAveragePooling2D()

        def forward_fn(di, is_training=True):
            return {'out': pool(di['in'])}

        return forward_fn

    return siso_tensorflow_eager_module('GlobalAveragePool', compile_fn, {}) def dropout(h_keep_prob):

    class Dropout(tf.keras.layers.Layer):

        def __init__(self, rate, seed=None, **kwargs):
            super(Dropout, self).__init__(**kwargs)
            self.rate = rate
            self.seed = seed
            self.supports_masking = True

        def call(self, inputs, training=None):
            if training is None:
                training = True

            def dropped_inputs():
                return tf.nn.dropout(inputs, 1 - self.rate, seed=self.seed)

            output = tf_utils.smart_cond(
                training, dropped_inputs, lambda: array_ops.identity(inputs))
            return output

        def compute_output_shape(self, input_shape):
            return input_shape

        def get_config(self):
            config = {'rate': self.rate, 'seed': self.seed}
            base_config = super(Dropout, self).get_config()
            return dict(list(base_config.items()) + list(config.items()))

    def compile_fn(di, dh):
        dropout_op = Dropout(1 - dh['keep_prob'])

        def forward_fn(di, is_training=True):
            # out = tf.nn.dropout(di['in'], dh['keep_prob'])
            # else:
            #     out = di['in']
            out = dropout_op(di['in'])
            return {'out': out}

        return forward_fn

    return siso_tensorflow_eager_module('Dropout', compile_fn,
                                        {'keep_prob': h_keep_prob}) def add(num_inputs):

    def compile_fn(di, dh):
        if num_inputs > 1:
            add = tf.keras.layers.Add()
        else:
            add = None

        def forward_fn(di, is_training=True):
            out = add([di[inp] for inp in di]) if add else di['in0']

            return {'out': out}

        return forward_fn

    return TensorflowEagerModule('Add', compile_fn, {},
                                 ['in' + str(i) for i in range(num_inputs)],
                                 ['out']).get_io() def flatten():

    def compile_fn(di, dh):
        flatten = tf.keras.layers.Flatten()

        def forward_fn(di, is_training=True):
            return {'out': flatten(di['in'])}

        return forward_fn

    return siso_tensorflow_eager_module('Flatten', compile_fn, {}) class Discrete(co.Hyperparameter):
    """List valued hyperparameter.

    This type of hyperparameter has a finite number of possible values.

    Args:
        vs (list[object]): List of possible parameter values that the
            hyperparameter can take.
        scope (deep_architect.core.Scope, optional): The scope in which to register the
            hyperparameter in.
        name (str, optional): Name from which the name of the hyperparameter
            in the scope is derived.
    """

    def __init__(self, vs, scope=None, name=None):
        assert len(vs) > 0
        co.Hyperparameter.__init__(self, scope, name)
        self.vs = vs

    def _check_value(self, val):
        """Checks if the chosen values is in the list of valid values.

        Asserts ``False`` if the value is not in the list.
        """
        if val not in self.vs:
            print(self.get_name())
            print(self.vs)
            print(val)
        assert val in self.vs
import tensorflow as tf
import deep_architect.core as co
import deep_architect.hyperparameters as hp
import deep_architect.helpers.tensorflow_eager_support as htfe
import deep_architect.modules as mo
from tensorflow.python.keras.utils import tf_utils
from deep_architect.contrib.deep_learning_backend.tensorflow_keras_ops import (
    relu, batch_normalization, conv2d, separable_conv2d, avg_pool2d, max_pool2d,
    min_pool2d, fc_layer, global_pool2d, dropout, add, flatten)
from deep_architect.hyperparameters import Discrete as D
    concat_ins, concat_outs = concat(num_ins)
    di = {}
    last_in, last_out = inputs[-1]
    for i in range(num_ins):
        f_in, f_out = inputs[i]
        factorized[i][0]['in0'].connect(f_out['out'])
        factorized[i][0]['in1'].connect(last_out['out'])
        di['in' + str(i)] = f_in['in']
        concat_ins['in' + str(i)].connect(factorized[i][1]['out'])

    return di, concat_outs


def maybe_factorized_reduction(add_relu=False):

    def compile_fn(di, dh):
        _, _, height, channels = di['in0'].shape.as_list()
        _, _, final_height, final_channels = di['in1'].shape.as_list()
        if add_relu:
            relu = tf.keras.layers.ReLU()
        if height == final_height and channels == final_channels:
            pass
        elif height == final_height:
            conv = tf.keras.layers.Conv2D(final_channels, 1)
            bn = tf.keras.layers.BatchNormalization()
        else:
            avg_pool = tf.keras.layers.AveragePooling2D(1, 2, padding='VALID')
            conv1 = tf.keras.layers.Conv2D(int(final_channels / 2), 1)
            conv2 = tf.keras.layers.Conv2D(int(final_channels / 2), 1)
            bn = tf.keras.layers.BatchNormalization()
            pad = tf.keras.layers.ZeroPadding2D(((0, 1), (0, 1)))
            slice_layer = tf.keras.layers.Lambda(lambda x: x[:, 1:, 1:, :])
            concat = tf.keras.layers.Concatenate(axis=3)

        def forward_fn(di, is_training=True):
            inp = relu(di['in0']) if add_relu else di['in0']
            if height == final_height and channels == final_channels:
                out = inp
            elif height == final_height:
                out = conv(inp)
                out = bn(out)
            else:
                path1 = avg_pool(inp)
                path1 = conv1(path1)
                path2 = slice_layer(pad(inp))
                path2 = avg_pool(path2)
                path2 = conv2(path2)
                out = concat([path1, path2])
                out = bn(out)
            return {'out': out}

        return forward_fn

    return htfe.TensorflowEagerModule('MaybeFactorizedReduction', compile_fn,
                                      {}, ['in0', 'in1'], ['out']).get_io()


def cell_input_fn(filters):
    prev_input = mo.identity()
    cur_input = wrap_relu_batch_norm(conv2d(D([filters]), D([1])))
    transformed_prev_input = maybe_factorized_reduction(add_relu=True)
    transformed_prev_input[0]['in0'].connect(prev_input[1]['out'])
    transformed_prev_input[0]['in1'].connect(cur_input[1]['out'])
    return {
        'in0': prev_input[0]['in'],
        'in1': cur_input[0]['in']
    }, {
        'out0': transformed_prev_input[1]['out'],
        'out1': cur_input[1]['out']
    }


def create_cell_generator(num_nodes, reduction):
    prefix = 'reduction' if reduction else 'normal'
    cell_ops = [
        SP1_ops('%s_op_%d_%d' % (prefix, i // 2, i % 2), reduction)
        for i in range(2 * num_nodes)
    ]
    connection_hparams = [
        D(list(range(i // 2 + 2)), name='%s_in_%d_%d' % (prefix, i // 2, i % 2))
        for i in range(2 * num_nodes)
    ]

    def generate(filters, cell_ratio):
        return cell(
            lambda: cell_input_fn(filters), lambda in_id, node_id, op_num:
            intermediate_node_fn(reduction, in_id, node_id, op_num, filters,
                                 cell_ratio, cell_ops), lambda: add(2),
            combine_unused, num_nodes, connection_hparams)

    return generate


def stem(filters):
    return mo.siso_sequential(
        [conv2d(D([filters]), D([3])),
         batch_normalization()])


def global_convolution(h_num_filters):

    def compile_fn(di, dh):
        _, h, w, _ = di['in'].shape.as_list()
        conv = tf.keras.layers.Conv2D(dh['num_filters'], [h, w],
                                      use_bias=False,
                                      padding='VALID')

        def forward_fn(di, is_training=True):
            return {'out': conv(di['in'])}

        return forward_fn

    return htfe.siso_tensorflow_eager_module('GlobalConv2D', compile_fn, {
        'num_filters': h_num_filters,
    })


def aux_logits():
    return mo.siso_sequential([
        relu(),

        avg_pool2d(D([5]), D([3]), D(['VALID'])),