# Path: deep_architect/contrib/misc/search_spaces/tensorflow/common.py
# D = hp.Discrete
# 
# def siso_tensorflow_module(name, compile_fn, name_to_hyperp, scope=None):
#     return htf.TensorflowModule(name, compile_fn, name_to_hyperp, ['in'],
#                                 ['out'], scope).get_io()


from deep_architect.contrib.misc.search_spaces.tensorflow.common import D, siso_tensorflow_module
import deep_architect.contrib.misc.search_spaces.tensorflow.dnn as dnn
import deep_architect.modules as mo
import tensorflow as tf
import numpy as np

        def forward_fn(di):
            return {
                'out':
                tf.nn.max_pool(di['in'],
                               [1, dh['kernel_size'], dh['kernel_size'], 1],
                               [1, dh['stride'], dh['stride'], 1], 'SAME')
            }

        return forward_fn

    return siso_tensorflow_module('MaxPool2D', compile_fn, {
        'kernel_size': h_kernel_size,
        'stride': h_stride,
    })


def conv_cell(h_num_filters, h_filter_width, h_swap, h_opt_drop, h_keep_prob,
              stride):
    assert stride >= 1
    return mo.siso_sequential([
        conv2d(h_num_filters, h_filter_width, D([stride]), D([0])),
        mo.siso_permutation([dnn.relu, dnn.batch_normalization], h_swap),
        mo.siso_optional(lambda: dnn.dropout(h_keep_prob), h_opt_drop)
    ])


def conv_net(h_num_spatial_reductions):
    h_swap = D([0, 1])
    h_opt_drop = D([0, 1])

# Path: deep_architect/contrib/misc/search_spaces/tensorflow/common.py
# D = hp.Discrete
# 
# def siso_tensorflow_module(name, compile_fn, name_to_hyperp, scope=None):
#     return htf.TensorflowModule(name, compile_fn, name_to_hyperp, ['in'],
#                                 ['out'], scope).get_io()


# Path: deep_architect/contrib/misc/search_spaces/tensorflow/cnn2d.py
from deep_architect.contrib.misc.search_spaces.tensorflow.common import D, siso_tensorflow_module
import deep_architect.contrib.misc.search_spaces.tensorflow.dnn as dnn
import deep_architect.modules as mo
import tensorflow as tf
import numpy as np


        def forward_fn(di):
            return {
                'out':
                tf.nn.max_pool(di['in'],
                               [1, dh['kernel_size'], dh['kernel_size'], 1],
                               [1, dh['stride'], dh['stride'], 1], 'SAME')
            }

        return forward_fn

    return siso_tensorflow_module('MaxPool2D', compile_fn, {
        'kernel_size': h_kernel_size,
        'stride': h_stride,
    })


def conv_cell(h_num_filters, h_filter_width, h_swap, h_opt_drop, h_keep_prob,
              stride):
    assert stride >= 1
    return mo.siso_sequential([
        conv2d(h_num_filters, h_filter_width, D([stride]), D([0])),
        mo.siso_permutation([dnn.relu, dnn.batch_normalization], h_swap),
        mo.siso_optional(lambda: dnn.dropout(h_keep_prob), h_opt_drop)
    ])


def conv_net(h_num_spatial_reductions):
    h_swap = D([0, 1])
    h_opt_drop = D([0, 1])

    h_keep_prob = D([0.25, 0.5, 0.75, 1.0])