--- a
+++ b/ext/neuron/models.py
@@ -0,0 +1,768 @@
+"""
+tensorflow/keras utilities for the neuron project
+
+If you use this code, please cite 
+Dalca AV, Guttag J, Sabuncu MR
+Anatomical Priors in Convolutional Networks for Unsupervised Biomedical Segmentation, 
+CVPR 2018
+
+Contact: adalca [at] csail [dot] mit [dot] edu
+License: GPLv3
+"""
+
+import sys
+
+from ext.neuron import layers
+
+# third party
+import numpy as np
+import tensorflow as tf
+import keras
+import keras.layers as KL
+from keras.models import Model
+import keras.backend as K
+
+
+def unet(nb_features,
+         input_shape,
+         nb_levels,
+         conv_size,
+         nb_labels,
+         name='unet',
+         prefix=None,
+         feat_mult=1,
+         pool_size=2,
+         use_logp=True,
+         padding='same',
+         dilation_rate_mult=1,
+         activation='elu',
+         skip_n_concatenations=0,
+         use_residuals=False,
+         final_pred_activation='softmax',
+         nb_conv_per_level=1,
+         add_prior_layer=False,
+         layer_nb_feats=None,
+         conv_dropout=0,
+         batch_norm=None,
+         input_model=None):
+    """
+    unet-style keras model with an overdose of parametrization.
+
+    Parameters:
+        nb_features: the number of features at each convolutional level
+            see below for `feat_mult` and `layer_nb_feats` for modifiers to this number
+        input_shape: input layer shape, vector of size ndims + 1 (nb_channels)
+        conv_size: the convolution kernel size
+        nb_levels: the number of Unet levels (number of downsamples) in the "encoder" 
+            (e.g. 4 would give you 4 levels in encoder, 4 in decoder)
+        nb_labels: number of output channels
+        name (default: 'unet'): the name of the network
+        prefix (default: `name` value): prefix to be added to layer names
+        feat_mult (default: 1) multiple for `nb_features` as we go down the encoder levels.
+            e.g. feat_mult of 2 and nb_features of 16 would yield 32 features in the 
+            second layer, 64 features in the third layer, etc.
+        pool_size (default: 2): max pooling size (integer or list if specifying per dimension)
+        skip_n_concatenations=0: enabled to skip concatenation links between contracting and expanding paths for the n
+            top levels.
+        use_logp:
+        padding:
+        dilation_rate_mult:
+        activation:
+        use_residuals:
+        final_pred_activation:
+        nb_conv_per_level:
+        add_prior_layer:
+        skip_n_concatenations:
+        layer_nb_feats: list of the number of features for each layer. Automatically used if specified
+        conv_dropout: dropout probability
+        batch_norm:
+        input_model: concatenate the provided input_model to this current model.
+            Only the first output of input_model is used.
+    """
+
+    # naming
+    model_name = name
+    if prefix is None:
+        prefix = model_name
+
+    # volume size data
+    ndims = len(input_shape) - 1
+    if isinstance(pool_size, int):
+        pool_size = (pool_size,) * ndims
+
+    # get encoding model
+    enc_model = conv_enc(nb_features,
+                         input_shape,
+                         nb_levels,
+                         conv_size,
+                         name=model_name,
+                         prefix=prefix,
+                         feat_mult=feat_mult,
+                         pool_size=pool_size,
+                         padding=padding,
+                         dilation_rate_mult=dilation_rate_mult,
+                         activation=activation,
+                         use_residuals=use_residuals,
+                         nb_conv_per_level=nb_conv_per_level,
+                         layer_nb_feats=layer_nb_feats,
+                         conv_dropout=conv_dropout,
+                         batch_norm=batch_norm,
+                         input_model=input_model)
+
+    # get decoder
+    # use_skip_connections=True makes it a u-net
+    lnf = layer_nb_feats[(nb_levels * nb_conv_per_level):] if layer_nb_feats is not None else None
+    dec_model = conv_dec(nb_features,
+                         [],
+                         nb_levels,
+                         conv_size,
+                         nb_labels,
+                         name=model_name,
+                         prefix=prefix,
+                         feat_mult=feat_mult,
+                         pool_size=pool_size,
+                         use_skip_connections=True,
+                         skip_n_concatenations=skip_n_concatenations,
+                         padding=padding,
+                         dilation_rate_mult=dilation_rate_mult,
+                         activation=activation,
+                         use_residuals=use_residuals,
+                         final_pred_activation='linear' if add_prior_layer else final_pred_activation,
+                         nb_conv_per_level=nb_conv_per_level,
+                         batch_norm=batch_norm,
+                         layer_nb_feats=lnf,
+                         conv_dropout=conv_dropout,
+                         input_model=enc_model)
+    final_model = dec_model
+
+    if add_prior_layer:
+        final_model = add_prior(dec_model,
+                                [*input_shape[:-1], nb_labels],
+                                name=model_name + '_prior',
+                                use_logp=use_logp,
+                                final_pred_activation=final_pred_activation)
+
+    return final_model
+
+
+def ae(nb_features,
+       input_shape,
+       nb_levels,
+       conv_size,
+       nb_labels,
+       enc_size,
+       name='ae',
+       feat_mult=1,
+       pool_size=2,
+       padding='same',
+       activation='elu',
+       use_residuals=False,
+       nb_conv_per_level=1,
+       batch_norm=None,
+       enc_batch_norm=None,
+       ae_type='conv',  # 'dense', or 'conv'
+       enc_lambda_layers=None,
+       add_prior_layer=False,
+       use_logp=True,
+       conv_dropout=0,
+       include_mu_shift_layer=False,
+       single_model=False,  # whether to return a single model, or a tuple of models that can be stacked.
+       final_pred_activation='softmax',
+       do_vae=False,
+       input_model=None):
+    """Convolutional Auto-Encoder. Optionally Variational (if do_vae is set to True)."""
+
+    # naming
+    model_name = name
+
+    # volume size data
+    ndims = len(input_shape) - 1
+    if isinstance(pool_size, int):
+        pool_size = (pool_size,) * ndims
+
+    # get encoding model
+    enc_model = conv_enc(nb_features,
+                         input_shape,
+                         nb_levels,
+                         conv_size,
+                         name=model_name,
+                         feat_mult=feat_mult,
+                         pool_size=pool_size,
+                         padding=padding,
+                         activation=activation,
+                         use_residuals=use_residuals,
+                         nb_conv_per_level=nb_conv_per_level,
+                         conv_dropout=conv_dropout,
+                         batch_norm=batch_norm,
+                         input_model=input_model)
+
+    # middle AE structure
+    if single_model:
+        in_input_shape = None
+        in_model = enc_model
+    else:
+        in_input_shape = enc_model.output.shape.as_list()[1:]
+        in_model = None
+    mid_ae_model = single_ae(enc_size,
+                             in_input_shape,
+                             conv_size=conv_size,
+                             name=model_name,
+                             ae_type=ae_type,
+                             input_model=in_model,
+                             batch_norm=enc_batch_norm,
+                             enc_lambda_layers=enc_lambda_layers,
+                             include_mu_shift_layer=include_mu_shift_layer,
+                             do_vae=do_vae)
+
+    # decoder
+    if single_model:
+        in_input_shape = None
+        in_model = mid_ae_model
+    else:
+        in_input_shape = mid_ae_model.output.shape.as_list()[1:]
+        in_model = None
+    dec_model = conv_dec(nb_features,
+                         in_input_shape,
+                         nb_levels,
+                         conv_size,
+                         nb_labels,
+                         name=model_name,
+                         feat_mult=feat_mult,
+                         pool_size=pool_size,
+                         use_skip_connections=False,
+                         padding=padding,
+                         activation=activation,
+                         use_residuals=use_residuals,
+                         final_pred_activation='linear',
+                         nb_conv_per_level=nb_conv_per_level,
+                         batch_norm=batch_norm,
+                         conv_dropout=conv_dropout,
+                         input_model=in_model)
+
+    if add_prior_layer:
+        dec_model = add_prior(dec_model,
+                              [*input_shape[:-1], nb_labels],
+                              name=model_name,
+                              prefix=model_name + '_prior',
+                              use_logp=use_logp,
+                              final_pred_activation=final_pred_activation)
+
+    if single_model:
+        return dec_model
+    else:
+        return dec_model, mid_ae_model, enc_model
+
+
+def conv_enc(nb_features,
+             input_shape,
+             nb_levels,
+             conv_size,
+             name=None,
+             prefix=None,
+             feat_mult=1,
+             pool_size=2,
+             dilation_rate_mult=1,
+             padding='same',
+             activation='elu',
+             layer_nb_feats=None,
+             use_residuals=False,
+             nb_conv_per_level=2,
+             conv_dropout=0,
+             batch_norm=None,
+             input_model=None):
+    """Fully Convolutional Encoder"""
+
+    # naming
+    model_name = name
+    if prefix is None:
+        prefix = model_name
+
+    # first layer: input
+    name = '%s_input' % prefix
+    if input_model is None:
+        input_tensor = KL.Input(shape=input_shape, name=name)
+        last_tensor = input_tensor
+    else:
+        input_tensor = input_model.inputs
+        last_tensor = input_model.outputs
+        if isinstance(last_tensor, list):
+            last_tensor = last_tensor[0]
+
+    # volume size data
+    ndims = len(input_shape) - 1
+    if isinstance(pool_size, int):
+        pool_size = (pool_size,) * ndims
+
+    # prepare layers
+    convL = getattr(KL, 'Conv%dD' % ndims)
+    conv_kwargs = {'padding': padding, 'activation': activation, 'data_format': 'channels_last'}
+    maxpool = getattr(KL, 'MaxPooling%dD' % ndims)
+
+    # down arm:
+    # add nb_levels of conv + ReLu + conv + ReLu. Pool after each of first nb_levels - 1 layers
+    lfidx = 0  # level feature index
+    for level in range(nb_levels):
+        lvl_first_tensor = last_tensor
+        nb_lvl_feats = np.round(nb_features * feat_mult ** level).astype(int)
+        conv_kwargs['dilation_rate'] = dilation_rate_mult ** level
+
+        for conv in range(nb_conv_per_level):  # does several conv per level, max pooling applied at the end
+            if layer_nb_feats is not None:  # None or List of all the feature numbers
+                nb_lvl_feats = layer_nb_feats[lfidx]
+                lfidx += 1
+
+            name = '%s_conv_downarm_%d_%d' % (prefix, level, conv)
+            if conv < (nb_conv_per_level - 1) or (not use_residuals):
+                last_tensor = convL(nb_lvl_feats, conv_size, **conv_kwargs, name=name)(last_tensor)
+            else:  # no activation
+                last_tensor = convL(nb_lvl_feats, conv_size, padding=padding, name=name)(last_tensor)
+
+            if conv_dropout > 0:
+                # conv dropout along feature space only
+                name = '%s_dropout_downarm_%d_%d' % (prefix, level, conv)
+                noise_shape = [None, *[1] * ndims, nb_lvl_feats]
+                last_tensor = KL.Dropout(conv_dropout, noise_shape=noise_shape, name=name)(last_tensor)
+
+        if use_residuals:
+            convarm_layer = last_tensor
+
+            # the "add" layer is the original input
+            # However, it may not have the right number of features to be added
+            nb_feats_in = lvl_first_tensor.get_shape()[-1]
+            nb_feats_out = convarm_layer.get_shape()[-1]
+            add_layer = lvl_first_tensor
+            if nb_feats_in > 1 and nb_feats_out > 1 and (nb_feats_in != nb_feats_out):
+                name = '%s_expand_down_merge_%d' % (prefix, level)
+                last_tensor = convL(nb_lvl_feats, conv_size, **conv_kwargs, name=name)(lvl_first_tensor)
+                add_layer = last_tensor
+
+                if conv_dropout > 0:
+                    noise_shape = [None, *[1] * ndims, nb_lvl_feats]
+                    convarm_layer = KL.Dropout(conv_dropout, noise_shape=noise_shape)(last_tensor)
+
+            name = '%s_res_down_merge_%d' % (prefix, level)
+            last_tensor = KL.add([add_layer, convarm_layer], name=name)
+
+            name = '%s_res_down_merge_act_%d' % (prefix, level)
+            last_tensor = KL.Activation(activation, name=name)(last_tensor)
+
+        if batch_norm is not None:
+            name = '%s_bn_down_%d' % (prefix, level)
+            last_tensor = KL.BatchNormalization(axis=batch_norm, name=name)(last_tensor)
+
+        # max pool if we're not at the last level
+        if level < (nb_levels - 1):
+            name = '%s_maxpool_%d' % (prefix, level)
+            last_tensor = maxpool(pool_size=pool_size, name=name, padding=padding)(last_tensor)
+
+    # create the model and return
+    model = Model(inputs=input_tensor, outputs=[last_tensor], name=model_name)
+    return model
+
+
+def conv_dec(nb_features,
+             input_shape,
+             nb_levels,
+             conv_size,
+             nb_labels,
+             name=None,
+             prefix=None,
+             feat_mult=1,
+             pool_size=2,
+             use_skip_connections=False,
+             skip_n_concatenations=0,
+             padding='same',
+             dilation_rate_mult=1,
+             activation='elu',
+             use_residuals=False,
+             final_pred_activation='softmax',
+             nb_conv_per_level=2,
+             layer_nb_feats=None,
+             batch_norm=None,
+             conv_dropout=0,
+             input_model=None):
+    """Fully Convolutional Decoder"""
+
+    # naming
+    model_name = name
+    if prefix is None:
+        prefix = model_name
+
+    # if using skip connections, make sure need to use them.
+    if use_skip_connections:
+        assert input_model is not None, "is using skip connections, tensors dictionary is required"
+
+    # first layer: input
+    input_name = '%s_input' % prefix
+    if input_model is None:
+        input_tensor = KL.Input(shape=input_shape, name=input_name)
+        last_tensor = input_tensor
+    else:
+        input_tensor = input_model.input
+        last_tensor = input_model.output
+        input_shape = last_tensor.shape.as_list()[1:]
+
+    # vol size info
+    ndims = len(input_shape) - 1
+    if isinstance(pool_size, int):
+        if ndims > 1:
+            pool_size = (pool_size,) * ndims
+
+    # prepare layers
+    convL = getattr(KL, 'Conv%dD' % ndims)
+    conv_kwargs = {'padding': padding, 'activation': activation}
+    upsample = getattr(KL, 'UpSampling%dD' % ndims)
+
+    # up arm:
+    # nb_levels - 1 layers of Deconvolution3D
+    #    (approx via up + conv + ReLu) + merge + conv + ReLu + conv + ReLu
+    lfidx = 0
+    for level in range(nb_levels - 1):
+        nb_lvl_feats = np.round(nb_features * feat_mult ** (nb_levels - 2 - level)).astype(int)
+        conv_kwargs['dilation_rate'] = dilation_rate_mult ** (nb_levels - 2 - level)
+
+        # upsample matching the max pooling layers size
+        name = '%s_up_%d' % (prefix, nb_levels + level)
+        last_tensor = upsample(size=pool_size, name=name)(last_tensor)
+        up_tensor = last_tensor
+
+        # merge layers combining previous layer
+        if use_skip_connections & (level < (nb_levels - skip_n_concatenations - 1)):
+            conv_name = '%s_conv_downarm_%d_%d' % (prefix, nb_levels - 2 - level, nb_conv_per_level - 1)
+            cat_tensor = input_model.get_layer(conv_name).output
+            name = '%s_merge_%d' % (prefix, nb_levels + level)
+            last_tensor = KL.concatenate([cat_tensor, last_tensor], axis=ndims + 1, name=name)
+
+        # convolution layers
+        for conv in range(nb_conv_per_level):
+            if layer_nb_feats is not None:
+                nb_lvl_feats = layer_nb_feats[lfidx]
+                lfidx += 1
+
+            name = '%s_conv_uparm_%d_%d' % (prefix, nb_levels + level, conv)
+            if conv < (nb_conv_per_level - 1) or (not use_residuals):
+                last_tensor = convL(nb_lvl_feats, conv_size, **conv_kwargs, name=name)(last_tensor)
+            else:
+                last_tensor = convL(nb_lvl_feats, conv_size, padding=padding, name=name)(last_tensor)
+
+            if conv_dropout > 0:
+                name = '%s_dropout_uparm_%d_%d' % (prefix, level, conv)
+                noise_shape = [None, *[1] * ndims, nb_lvl_feats]
+                last_tensor = KL.Dropout(conv_dropout, noise_shape=noise_shape, name=name)(last_tensor)
+
+        # residual block
+        if use_residuals:
+
+            # the "add" layer is the original input
+            # However, it may not have the right number of features to be added
+            add_layer = up_tensor
+            nb_feats_in = add_layer.get_shape()[-1]
+            nb_feats_out = last_tensor.get_shape()[-1]
+            if nb_feats_in > 1 and nb_feats_out > 1 and (nb_feats_in != nb_feats_out):
+                name = '%s_expand_up_merge_%d' % (prefix, level)
+                add_layer = convL(nb_lvl_feats, conv_size, **conv_kwargs, name=name)(add_layer)
+
+                if conv_dropout > 0:
+                    noise_shape = [None, *[1] * ndims, nb_lvl_feats]
+                    last_tensor = KL.Dropout(conv_dropout, noise_shape=noise_shape)(last_tensor)
+
+            name = '%s_res_up_merge_%d' % (prefix, level)
+            last_tensor = KL.add([last_tensor, add_layer], name=name)
+
+            name = '%s_res_up_merge_act_%d' % (prefix, level)
+            last_tensor = KL.Activation(activation, name=name)(last_tensor)
+
+        if batch_norm is not None:
+            name = '%s_bn_up_%d' % (prefix, level)
+            last_tensor = KL.BatchNormalization(axis=batch_norm, name=name)(last_tensor)
+
+    # Compute likelihood prediction (no activation yet)
+    name = '%s_likelihood' % prefix
+    last_tensor = convL(nb_labels, 1, activation=None, name=name)(last_tensor)
+    like_tensor = last_tensor
+
+    # output prediction layer
+    # we use a softmax to compute P(L_x|I) where x is each location
+    if final_pred_activation == 'softmax':
+        name = '%s_prediction' % prefix
+        softmax_lambda_fcn = lambda x: keras.activations.softmax(x, axis=ndims + 1)
+        pred_tensor = KL.Lambda(softmax_lambda_fcn, name=name)(last_tensor)
+
+    # otherwise create a layer that does nothing.
+    else:
+        name = '%s_prediction' % prefix
+        pred_tensor = KL.Activation('linear', name=name)(like_tensor)
+
+    # create the model and return
+    model = Model(inputs=input_tensor, outputs=pred_tensor, name=model_name)
+    return model
+
+
+def add_prior(input_model,
+              prior_shape,
+              name='prior_model',
+              prefix=None,
+              use_logp=True,
+              final_pred_activation='softmax'):
+    """
+    Append post-prior layer to a given model
+    """
+
+    # naming
+    model_name = name
+    if prefix is None:
+        prefix = model_name
+
+    # prior input layer
+    prior_input_name = '%s-input' % prefix
+    prior_tensor = KL.Input(shape=prior_shape, name=prior_input_name)
+    prior_tensor_input = prior_tensor
+    like_tensor = input_model.output
+
+    # operation varies depending on whether we log() prior or not.
+    if use_logp:
+        print("Breaking change: use_logp option now requires log input!", file=sys.stderr)
+        merge_op = KL.add
+
+    else:
+        # using sigmoid to get the likelihood values between 0 and 1
+        # note: they won't add up to 1.
+        name = '%s_likelihood_sigmoid' % prefix
+        like_tensor = KL.Activation('sigmoid', name=name)(like_tensor)
+        merge_op = KL.multiply
+
+    # merge the likelihood and prior layers into posterior layer
+    name = '%s_posterior' % prefix
+    post_tensor = merge_op([prior_tensor, like_tensor], name=name)
+
+    # output prediction layer
+    # we use a softmax to compute P(L_x|I) where x is each location
+    pred_name = '%s_prediction' % prefix
+    if final_pred_activation == 'softmax':
+        assert use_logp, 'cannot do softmax when adding prior via P()'
+        print("using final_pred_activation %s for %s" % (final_pred_activation, model_name))
+        softmax_lambda_fcn = lambda x: keras.activations.softmax(x, axis=-1)
+        pred_tensor = KL.Lambda(softmax_lambda_fcn, name=pred_name)(post_tensor)
+
+    else:
+        pred_tensor = KL.Activation('linear', name=pred_name)(post_tensor)
+
+    # create the model
+    model_inputs = [*input_model.inputs, prior_tensor_input]
+    model = Model(inputs=model_inputs, outputs=[pred_tensor], name=model_name)
+
+    # compile
+    return model
+
+
+def single_ae(enc_size,
+              input_shape,
+              name='single_ae',
+              prefix=None,
+              ae_type='dense',  # 'dense', or 'conv'
+              conv_size=None,
+              input_model=None,
+              enc_lambda_layers=None,
+              batch_norm=True,
+              padding='same',
+              activation=None,
+              include_mu_shift_layer=False,
+              do_vae=False):
+    """single-layer Autoencoder (i.e. input - encoding - output"""
+
+    # naming
+    model_name = name
+    if prefix is None:
+        prefix = model_name
+
+    if enc_lambda_layers is None:
+        enc_lambda_layers = []
+
+    # prepare input
+    input_name = '%s_input' % prefix
+    if input_model is None:
+        assert input_shape is not None, 'input_shape of input_model is necessary'
+        input_tensor = KL.Input(shape=input_shape, name=input_name)
+        last_tensor = input_tensor
+    else:
+        input_tensor = input_model.input
+        last_tensor = input_model.output
+        input_shape = last_tensor.shape.as_list()[1:]
+    input_nb_feats = last_tensor.shape.as_list()[-1]
+
+    # prepare conv type based on input
+    ndims = len(input_shape) - 1
+    if ae_type == 'conv':
+        convL = getattr(KL, 'Conv%dD' % ndims)
+        assert conv_size is not None, 'with conv ae, need conv_size'
+        conv_kwargs = {'padding': padding, 'activation': activation}
+        enc_size_str = None
+
+    # if want to go through a dense layer in the middle of the U, need to:
+    # - flatten last layer if not flat
+    # - do dense encoding and decoding
+    # - unflatten (reshape spatially) at end
+    else:  # ae_type == 'dense'
+        if len(input_shape) > 1:
+            name = '%s_ae_%s_down_flat' % (prefix, ae_type)
+            last_tensor = KL.Flatten(name=name)(last_tensor)
+        convL = conv_kwargs = None
+        assert len(enc_size) == 1, "enc_size should be of length 1 for dense layer"
+        enc_size_str = ''.join(['%d_' % d for d in enc_size])[:-1]
+
+    # recall this layer
+    pre_enc_layer = last_tensor
+
+    # encoding layer
+    if ae_type == 'dense':
+        name = '%s_ae_mu_enc_dense_%s' % (prefix, enc_size_str)
+        last_tensor = KL.Dense(enc_size[0], name=name)(pre_enc_layer)
+
+    else:  # convolution
+
+        # convolve then resize. enc_size should be [nb_dim1, nb_dim2, ..., nb_feats]
+        assert len(enc_size) == len(input_shape), \
+            "encoding size does not match input shape %d %d" % (len(enc_size), len(input_shape))
+
+        if list(enc_size)[:-1] != list(input_shape)[:-1] and \
+                all([f is not None for f in input_shape[:-1]]) and \
+                all([f is not None for f in enc_size[:-1]]):
+
+            name = '%s_ae_mu_enc_conv' % prefix
+            last_tensor = convL(enc_size[-1], conv_size, name=name, **conv_kwargs)(pre_enc_layer)
+
+            name = '%s_ae_mu_enc' % prefix
+            zf = [enc_size[:-1][f] / last_tensor.shape.as_list()[1:-1][f] for f in range(len(enc_size) - 1)]
+            last_tensor = layers.Resize(zoom_factor=zf, name=name)(last_tensor)
+
+        elif enc_size[-1] is None:  # convolutional, but won't tell us bottleneck
+            name = '%s_ae_mu_enc' % prefix
+            last_tensor = KL.Lambda(lambda x: x, name=name)(pre_enc_layer)
+
+        else:
+            name = '%s_ae_mu_enc' % prefix
+            last_tensor = convL(enc_size[-1], conv_size, name=name, **conv_kwargs)(pre_enc_layer)
+
+    if include_mu_shift_layer:
+        # shift
+        name = '%s_ae_mu_shift' % prefix
+        last_tensor = layers.LocalBias(name=name)(last_tensor)
+
+    # encoding clean-up layers
+    for layer_fcn in enc_lambda_layers:
+        lambda_name = layer_fcn.__name__
+        name = '%s_ae_mu_%s' % (prefix, lambda_name)
+        last_tensor = KL.Lambda(layer_fcn, name=name)(last_tensor)
+
+    if batch_norm is not None:
+        name = '%s_ae_mu_bn' % prefix
+        last_tensor = KL.BatchNormalization(axis=batch_norm, name=name)(last_tensor)
+
+    # have a simple layer that does nothing to have a clear name before sampling
+    name = '%s_ae_mu' % prefix
+    last_tensor = KL.Lambda(lambda x: x, name=name)(last_tensor)
+
+    # if doing variational AE, will need the sigma layer as well.
+    if do_vae:
+        mu_tensor = last_tensor
+
+        # encoding layer
+        if ae_type == 'dense':
+            name = '%s_ae_sigma_enc_dense_%s' % (prefix, enc_size_str)
+            last_tensor = KL.Dense(enc_size[0], name=name)(pre_enc_layer)
+
+        else:
+            if list(enc_size)[:-1] != list(input_shape)[:-1] and \
+                    all([f is not None for f in input_shape[:-1]]) and \
+                    all([f is not None for f in enc_size[:-1]]):
+
+                assert len(enc_size) - 1 == 2, "Sorry, I have not yet implemented non-2D resizing..."
+                name = '%s_ae_sigma_enc_conv' % prefix
+                last_tensor = convL(enc_size[-1], conv_size, name=name, **conv_kwargs)(pre_enc_layer)
+
+                name = '%s_ae_sigma_enc' % prefix
+                resize_fn = lambda x: tf.image.resize_bilinear(x, enc_size[:-1])
+                last_tensor = KL.Lambda(resize_fn, name=name)(last_tensor)
+
+            elif enc_size[-1] is None:  # convolutional, but won't tell us bottleneck
+                name = '%s_ae_sigma_enc' % prefix
+                last_tensor = convL(pre_enc_layer.shape.as_list()[-1], conv_size, name=name, **conv_kwargs)(
+                    pre_enc_layer)
+                # cannot use lambda, then mu and sigma will be same layer.
+                # last_tensor = KL.Lambda(lambda x: x, name=name)(pre_enc_layer)
+
+            else:
+                name = '%s_ae_sigma_enc' % prefix
+                last_tensor = convL(enc_size[-1], conv_size, name=name, **conv_kwargs)(pre_enc_layer)
+
+        # encoding clean-up layers
+        for layer_fcn in enc_lambda_layers:
+            lambda_name = layer_fcn.__name__
+            name = '%s_ae_sigma_%s' % (prefix, lambda_name)
+            last_tensor = KL.Lambda(layer_fcn, name=name)(last_tensor)
+
+        if batch_norm is not None:
+            name = '%s_ae_sigma_bn' % prefix
+            last_tensor = KL.BatchNormalization(axis=batch_norm, name=name)(last_tensor)
+
+        # have a simple layer that does nothing to have a clear name before sampling
+        name = '%s_ae_sigma' % prefix
+        last_tensor = KL.Lambda(lambda x: x, name=name)(last_tensor)
+
+        logvar_tensor = last_tensor
+
+        # VAE sampling
+        sampler = _VAESample().sample_z
+
+        name = '%s_ae_sample' % prefix
+        last_tensor = KL.Lambda(sampler, name=name)([mu_tensor, logvar_tensor])
+
+    if include_mu_shift_layer:
+        # shift
+        name = '%s_ae_sample_shift' % prefix
+        last_tensor = layers.LocalBias(name=name)(last_tensor)
+
+    # decoding layer
+    if ae_type == 'dense':
+        name = '%s_ae_%s_dec_flat_%s' % (prefix, ae_type, enc_size_str)
+        last_tensor = KL.Dense(np.prod(input_shape), name=name)(last_tensor)
+
+        # unflatten if dense method
+        if len(input_shape) > 1:
+            name = '%s_ae_%s_dec' % (prefix, ae_type)
+            last_tensor = KL.Reshape(input_shape, name=name)(last_tensor)
+
+    else:
+
+        if list(enc_size)[:-1] != list(input_shape)[:-1] and \
+                all([f is not None for f in input_shape[:-1]]) and \
+                all([f is not None for f in enc_size[:-1]]):
+            name = '%s_ae_mu_dec' % prefix
+            zf = [last_tensor.shape.as_list()[1:-1][f] / enc_size[:-1][f] for f in range(len(enc_size) - 1)]
+            last_tensor = layers.Resize(zoom_factor=zf, name=name)(last_tensor)
+
+        name = '%s_ae_%s_dec' % (prefix, ae_type)
+        last_tensor = convL(input_nb_feats, conv_size, name=name, **conv_kwargs)(last_tensor)
+
+    if batch_norm is not None:
+        name = '%s_bn_ae_%s_dec' % (prefix, ae_type)
+        last_tensor = KL.BatchNormalization(axis=batch_norm, name=name)(last_tensor)
+
+    # create the model and return
+    model = Model(inputs=input_tensor, outputs=[last_tensor], name=model_name)
+    return model
+
+
+###############################################################################
+# Helper function
+###############################################################################
+
+class _VAESample:
+    def __init__(self):
+        pass
+
+    def sample_z(self, args):
+        mu, log_var = args
+        shape = K.shape(mu)
+        eps = K.random_normal(shape=shape, mean=0., stddev=1.)
+        return mu + K.exp(log_var / 2) * eps