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
            # Transform hidden states
            h0 = node_fn(idx0, i, 0)
            h1 = node_fn(idx1, i, 1)
            h0[0]['in'].connect(nodes[idx0])
            h1[0]['in'].connect(nodes[idx1])
            used_node[idx0] = used_node[idx1] = True

            # Combine hidden states
            h = combine_fn()
            h[0]['in0'].connect(h0[1]['out'])
            h[0]['in1'].connect(h1[1]['out'])

            nodes.append(h[1]['out'])

        ins, outs = unused_combine_fn(sum(not used for used in used_node))
        input_id = 0
        for ix, node in enumerate(nodes):
            if not used_node[ix]:
                ins['in' + str(input_id)].connect(node)
                input_id += 1
        return c_ins, outs

    name_to_hyperp = {
        '%d_%d' % (i // 2, i % 2): hyperparameters[i]
        for i in range(len(hyperparameters))
    }
    return mo.substitution_module('Cell', substitution_fn, name_to_hyperp,
                                  ['in0', 'in1'], ['out'], None)


def SP1_ops(name=None, reduction=False):
    if reduction:
        # Dilated convolution can't be done with strides
        ops = [
            'none', 'depth_sep3', 'depth_sep5', 'depth_sep7', 'avg3', 'max3',
            '1x7_7x1'
        ]
    else:
        ops = [
            'none', 'depth_sep3', 'depth_sep5', 'depth_sep7', 'avg3', 'max3',
            'dilated_3x3_rate_2', '1x7_7x1'
        ]
    return D(ops, name=name)


def wrap_relu_batch_norm(op):
    return mo.siso_sequential([relu(), op, batch_normalization()])


def conv_op(filters, filter_size, stride, dilation_rate, spatial_separable):
    if spatial_separable:
        return mo.siso_sequential([
            conv2d(D([filters]), D([[1, filter_size]]), D([[1, stride]])),
            batch_normalization(),
            relu(),
            conv2d(D([filters]), D([[filter_size, 1]]), D([[stride, 1]])),
        ])
    else:
        return conv2d(D([filters]), D([filter_size]), D([stride]),
                      D([dilation_rate]))


def full_conv_op(filters, filter_size, stride, dilation_rate,
                 spatial_separable):
    # Add bottleneck layer according to
    # https://github.com/tensorflow/tpu/blob/master/models/official/amoeba_net/network_utils.py
    if filter_size == 3 and spatial_separable:
        reduced_filter_size = int(3 * filters / 8)
    else:
        reduced_filter_size = int(filters / 4)
    if reduced_filter_size < 1:
        return wrap_relu_batch_norm(
            conv_op(filters, filter_size, stride, dilation_rate,
                    spatial_separable))
    else:
        return mo.siso_sequential([
            wrap_relu_batch_norm(conv2d(D([reduced_filter_size]), D([1]))),
            wrap_relu_batch_norm(
                conv_op(reduced_filter_size, filter_size, stride, dilation_rate,
                        spatial_separable)),
            wrap_relu_batch_norm(conv2d(D([filters]), D([1])))
        ])


def check_filters(filters, stride=1):

    def compile_fn(di, dh):
        num_filters = di['in'].shape[-1].value
        if num_filters != filters or stride > 1:
            conv = tf.keras.layers.Conv2D(filters,
                                          1,
                                          strides=stride,
                                          padding='SAME')
        else:
            conv = None

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

        return forward_fn

    return htfe.siso_tensorflow_eager_module('CheckFilters', compile_fn, {})


def pool_op(filters, filter_size, stride, pool_type):
    if pool_type == 'avg':
        pool = avg_pool2d(D([filter_size]), D([stride]))
    elif pool_type == 'max':
        pool = max_pool2d(D([filter_size]), D([stride]))
    else:
        pool = min_pool2d(D([filter_size]), D([stride]))

    return mo.siso_sequential([pool, check_filters(filters)])


def separable_conv_op(filters, filter_size, stride):
    return mo.siso_sequential([
        wrap_relu_batch_norm(
            separable_conv2d(D([filters]), D([filter_size]), D([stride]))),
        wrap_relu_batch_norm(

            separable_conv2d(D([filters]), D([filter_size]), D([1])))