# 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

        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'])}


# 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


        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