Upload ops.py

ops.py

@@ -0,0 +1,1067 @@
+from __future__ import division
+import math
+import tensorflow as tf
+
+from mi_gru_cell import MiGRUCell
+from mi_lstm_cell import MiLSTMCell
+from config import config
+
+eps = 1e-20
+inf = 1e30
+
+####################################### variables ########################################
+
+'''
+Initializes a weight matrix variable given a shape and a name. 
+Uses random_normal initialization if 1d, otherwise uses xavier. 
+'''
+def getWeight(shape, name = ""):
+    with tf.variable_scope("weights"):               
+        initializer = tf.contrib.layers.xavier_initializer()
+        # if len(shape) == 1: # good?
+        #     initializer = tf.random_normal_initializer()        
+        W = tf.get_variable("weight" + name, shape = shape, initializer = initializer)
+    return W
+
+'''
+Initializes a weight matrix variable given a shape and a name. Uses xavier
+'''
+def getKernel(shape, name = ""):
+    with tf.variable_scope("kernels"):               
+        initializer = tf.contrib.layers.xavier_initializer()
+        W = tf.get_variable("kernel" + name, shape = shape, initializer = initializer)
+    return W
+
+'''
+Initializes a bias variable given a shape and a name.
+'''
+def getBias(shape, name = ""):
+    with tf.variable_scope("biases"):              
+        initializer = tf.zeros_initializer()
+        b = tf.get_variable("bias" + name, shape = shape, initializer = initializer)
+    return b
+
+######################################### basics #########################################
+
+'''
+Multiplies input inp of any depth by a 2d weight matrix.  
+'''
+# switch with conv 1?
+def multiply(inp, W):
+    inDim = tf.shape(W)[0]
+    outDim = tf.shape(W)[1] 
+    newDims = tf.concat([tf.shape(inp)[:-1], tf.fill((1,), outDim)], axis = 0)
+    
+    inp = tf.reshape(inp, (-1, inDim))
+    output = tf.matmul(inp, W)
+    output = tf.reshape(output, newDims)
+
+    return output
+
+'''
+Concatenates x and y. Support broadcasting. 
+Optionally concatenate multiplication of x * y
+'''
+def concat(x, y, dim, mul = False, extendY = False):
+    if extendY:
+        y = tf.expand_dims(y, axis = -2)
+        # broadcasting to have the same shape
+        y = tf.zeros_like(x) + y
+
+    if mul:
+        out = tf.concat([x, y, x * y], axis = -1)
+        dim *= 3
+    else:
+        out = tf.concat([x, y], axis = -1)
+        dim *= 2
+    
+    return out, dim
+
+'''
+Adds L2 regularization for weight and kernel variables.
+'''
+# add l2 in the tf way
+def L2RegularizationOp(l2 = None):
+    if l2 is None:
+        l2 = config.l2
+    l2Loss = 0
+    names = ["weight", "kernel"]
+    for var in tf.trainable_variables():
+        if any((name in var.name.lower()) for name in names):
+            l2Loss += tf.nn.l2_loss(var)
+    return l2 * l2Loss
+
+######################################### attention #########################################
+
+'''
+Transform vectors to scalar logits.
+
+Args:
+    interactions: input vectors
+    [batchSize, N, dim]
+
+    dim: dimension of input vectors
+
+    sumMod: LIN for linear transformation to scalars.
+            SUM to sum up vectors entries to get scalar logit.
+
+    dropout: dropout value over inputs (for linear case)
+
+Return matching scalar for each interaction.
+[batchSize, N]
+'''
+sumMod = ["LIN", "SUM"]
+def inter2logits(interactions, dim, sumMod = "LIN", dropout = 1.0, name = "", reuse = None):
+    with tf.variable_scope("inter2logits" + name, reuse = reuse): 
+        if sumMod == "SUM":
+            logits = tf.reduce_sum(interactions, axis = -1)
+        else: # "LIN"
+            logits = linear(interactions, dim, 1, dropout = dropout, name = "logits")
+    return logits
+
+'''
+Transforms vectors to probability distribution. 
+Calls inter2logits and then softmax over these.
+
+Args:
+    interactions: input vectors
+    [batchSize, N, dim]
+
+    dim: dimension of input vectors
+
+    sumMod: LIN for linear transformation to scalars.
+            SUM to sum up vectors entries to get scalar logit.
+
+    dropout: dropout value over inputs (for linear case)
+
+Return attention distribution over interactions.
+[batchSize, N]
+'''
+def inter2att(interactions, dim, dropout = 1.0, name = "", reuse = None):
+    with tf.variable_scope("inter2att" + name, reuse = reuse): 
+        logits = inter2logits(interactions, dim, dropout = dropout)
+        attention = tf.nn.softmax(logits)    
+    return attention
+
+'''
+Sums up features using attention distribution to get a weighted average over them. 
+'''
+def att2Smry(attention, features):
+    return tf.reduce_sum(tf.expand_dims(attention, axis = -1) * features, axis = -2)
+
+####################################### activations ########################################
+
+'''
+Performs a variant of ReLU based on config.relu
+    PRM for PReLU
+    ELU for ELU
+    LKY for Leaky ReLU
+    otherwise, standard ReLU
+'''
+def relu(inp):                  
+    if config.relu == "PRM":
+        with tf.variable_scope(None, default_name = "prelu"):
+            alpha = tf.get_variable("alpha", shape = inp.get_shape()[-1], 
+                initializer = tf.constant_initializer(0.25))
+            pos = tf.nn.relu(inp)
+            neg = - (alpha * tf.nn.relu(-inp))
+            output = pos + neg
+    elif config.relu == "ELU":
+        output = tf.nn.elu(inp)
+    # elif config.relu == "SELU":
+    #     output = tf.nn.selu(inp) 
+    elif config.relu == "LKY":
+        # output = tf.nn.leaky_relu(inp, config.reluAlpha)
+        output = tf.maximum(inp, config.reluAlpha * inp)
+    elif config.relu == "STD": # STD
+        output = tf.nn.relu(inp)
+    
+    return output
+
+activations = {
+    "NON":      tf.identity, # lambda inp: inp    
+    "TANH":     tf.tanh,
+    "SIGMOID":  tf.sigmoid,
+    "RELU":     relu,
+    "ELU":      tf.nn.elu
+}    
+
+# Sample from Gumbel(0, 1)
+def sampleGumbel(shape): 
+    U = tf.random_uniform(shape, minval = 0, maxval = 1)
+    return -tf.log(-tf.log(U + eps) + eps)
+
+# Draw a clevr_sample from the Gumbel-Softmax distribution
+def gumbelSoftmaxSample(logits, temperature): 
+    y = logits + sampleGumbel(tf.shape(logits))
+    return tf.nn.softmax(y / temperature)
+
+def gumbelSoftmax(logits, temperature, train): # hard = False
+    # Sample from the Gumbel-Softmax distribution and optionally discretize.
+    # Args:
+    #    logits: [batch_size, n_class] unnormalized log-probs
+    #    temperature: non-negative scalar
+    #    hard: if True, take argmax, but differentiate w.r.t. soft clevr_sample y
+    # Returns:
+    #    [batch_size, n_class] clevr_sample from the Gumbel-Softmax distribution.
+    #    If hard=True, then the returned clevr_sample will be one-hot, otherwise it will
+    #    be a probabilitiy distribution that sums to 1 across classes
+
+    y = gumbelSoftmaxSample(logits, temperature)
+
+    # k = tf.shape(logits)[-1]
+    # yHard = tf.cast(tf.one_hot(tf.argmax(y,1),k), y.dtype)
+    yHard = tf.cast(tf.equal(y, tf.reduce_max(y, 1, keep_dims = True)), y.dtype)
+    yNew = tf.stop_gradient(yHard - y) + y
+
+    if config.gumbelSoftmaxBoth:
+        return y
+    if config.gumbelArgmaxBoth:
+        return yNew
+    ret = tf.cond(train, lambda: y, lambda: yNew)
+    
+    return ret 
+
+def softmaxDiscrete(logits, temperature, train):
+    if config.gumbelSoftmax:
+        return gumbelSoftmax(logits, temperature = temperature, train = train)
+    else:
+        return tf.nn.softmax(logits)
+
+def parametricDropout(name, train):
+    var = tf.get_variable("varDp" + name, shape = (), initializer = tf.constant_initializer(2), 
+        dtype = tf.float32)
+    dropout = tf.cond(train, lambda: tf.sigmoid(var), lambda: 1.0)
+    return dropout
+
+###################################### sequence helpers ######################################
+
+'''
+Casts exponential mask over a sequence with sequence length.
+Used to prepare logits before softmax.
+'''
+def expMask(seq, seqLength):
+    maxLength = tf.shape(seq)[-1]
+    mask = (1 - tf.cast(tf.sequence_mask(seqLength, maxLength), tf.float32)) * (-inf)
+    masked = seq + mask
+    return masked
+
+'''
+Computes seq2seq loss between logits and target sequences, with given lengths.
+'''
+def seq2SeqLoss(logits, targets, lengths):
+    mask = tf.sequence_mask(lengths, maxlen = tf.shape(targets)[1])
+    loss = tf.contrib.seq2seq.sequence_loss(logits, targets, tf.to_float(mask))
+    return loss
+
+'''
+Computes seq2seq loss between logits and target sequences, with given lengths.
+    acc1: accuracy per symbol 
+    acc2: accuracy per sequence
+'''
+def seq2seqAcc(preds, targets, lengths):
+    mask = tf.sequence_mask(lengths, maxlen = tf.shape(targets)[1])
+    corrects = tf.logical_and(tf.equal(preds, targets), mask)
+    numCorrects = tf.reduce_sum(tf.to_int32(corrects), axis = 1)
+    
+    acc1 = tf.to_float(numCorrects) / (tf.to_float(lengths) + eps) # add small eps instead?
+    acc1 = tf.reduce_mean(acc1)  
+    
+    acc2 = tf.to_float(tf.equal(numCorrects, lengths))
+    acc2 = tf.reduce_mean(acc2)      
+
+    return acc1, acc2
+
+########################################### linear ###########################################
+
+'''
+linear transformation.
+
+Args:
+    inp: input to transform
+    inDim: input dimension
+    outDim: output dimension
+    dropout: dropout over input
+    batchNorm: if not None, applies batch normalization to inputs
+    addBias: True to add bias
+    bias: initial bias value
+    act: if not None, activation to use after linear transformation
+    actLayer: if True and act is not None, applies another linear transformation on top of previous
+    actDropout: dropout to apply in the optional second linear transformation
+    retVars: if True, return parameters (weight and bias) 
+
+Returns linear transformation result.
+'''
+# batchNorm = {"decay": float, "train": Tensor}
+# actLayer: if activation is not non, stack another linear layer
+# maybe change naming scheme such that if name = "" than use it as default_name (-->unique?)
+def linear(inp, inDim, outDim, dropout = 1.0, 
+    batchNorm = None, addBias = True, bias = 0.0,
+    act = "NON", actLayer = True, actDropout = 1.0, 
+    retVars = False, name = "", reuse = None):
+    
+    with tf.variable_scope("linearLayer" + name, reuse = reuse):        
+        W = getWeight((inDim, outDim) if outDim > 1 else (inDim, ))
+        b = getBias((outDim, ) if outDim > 1 else ()) + bias
+        
+        if batchNorm is not None:
+            inp = tf.contrib.layers.batch_norm(inp, decay = batchNorm["decay"], 
+                center = True, scale = True, is_training = batchNorm["train"], updates_collections = None)
+            # tf.layers.batch_normalization, axis -1 ?
+
+        inp = tf.nn.dropout(inp, dropout)                
+        
+        if outDim > 1:
+            output = multiply(inp, W)
+        else:
+            output = tf.reduce_sum(inp * W, axis = -1)
+        
+        if addBias:
+            output += b
+
+        output = activations[act](output)
+
+        # good?
+        if act != "NON" and actLayer:
+            output = linear(output, outDim, outDim, dropout = actDropout, batchNorm = batchNorm,  
+                addBias = addBias, act = "NON", actLayer = False, 
+                name = name + "_2", reuse = reuse)
+
+    if retVars:
+        return (output, (W, b))
+
+    return output
+
+'''
+Computes Multi-layer feed-forward network.
+
+Args:
+    features: input features
+    dims: list with dimensions of network. 
+          First dimension is of the inputs, final is of the outputs.
+    batchNorm: if not None, applies batchNorm
+    dropout: dropout value to apply for each layer
+    act: activation to apply between layers.
+    NON, TANH, SIGMOID, RELU, ELU
+'''
+# no activation after last layer
+# batchNorm = {"decay": float, "train": Tensor}
+def FCLayer(features, dims, batchNorm = None, dropout = 1.0, act = "RELU"):
+    layersNum = len(dims) - 1
+    
+    for i in range(layersNum):
+        features = linear(features, dims[i], dims[i+1], name = "fc_%d" % i, 
+            batchNorm = batchNorm, dropout = dropout)
+        # not the last layer
+        if i < layersNum - 1: 
+            features = activations[act](features)
+    
+    return features   
+
+###################################### cnns ######################################
+
+'''
+Computes convolution.
+
+Args:
+    inp: input features
+    inDim: input dimension
+    outDim: output dimension
+    batchNorm: if not None, applies batchNorm on inputs
+    dropout: dropout value to apply on inputs
+    addBias: True to add bias
+    kernelSize: kernel size
+    stride: stride size
+    act: activation to apply on outputs
+    NON, TANH, SIGMOID, RELU, ELU
+'''
+# batchNorm = {"decay": float, "train": Tensor, "center": bool, "scale": bool}
+# collections.namedtuple("batchNorm", ("decay", "train"))
+def cnn(inp, inDim, outDim, batchNorm = None, dropout = 1.0, addBias = True, 
+    kernelSize = None, stride = 1, act = "NON", name = "", reuse = None):
+    
+    with tf.variable_scope("cnnLayer" + name, reuse = reuse):
+        
+        if kernelSize is None:
+            kernelSize = config.stemKernelSize            
+        kernelH = kernelW = kernelSize
+        
+        kernel = getKernel((kernelH, kernelW, inDim, outDim))
+        b = getBias((outDim, ))
+        
+        if batchNorm is not None:
+            inp = tf.contrib.layers.batch_norm(inp, decay = batchNorm["decay"], center = batchNorm["center"], 
+                scale = batchNorm["scale"], is_training = batchNorm["train"], updates_collections = None)   
+
+        inp = tf.nn.dropout(inp, dropout)                
+        
+        output = tf.nn.conv2d(inp, filter = kernel, strides = [1, stride, stride, 1], padding = "SAME")
+        
+        if addBias:
+            output += b
+
+        output = activations[act](output)
+
+    return output
+
+'''
+Computes Multi-layer convolutional network.
+
+Args:
+    features: input features
+    dims: list with dimensions of network. 
+          First dimension is of the inputs. Final is of the outputs.
+    batchNorm: if not None, applies batchNorm
+    dropout: dropout value to apply for each layer
+    kernelSizes: list of kernel sizes for each layer. Default to config.stemKernelSize
+    strides: list of strides for each layer. Default to 1.
+    act: activation to apply between layers.
+    NON, TANH, SIGMOID, RELU, ELU
+'''
+# batchNorm = {"decay": float, "train": Tensor, "center": bool, "scale": bool}
+# activation after last layer
+def CNNLayer(features, dims, batchNorm = None, dropout = 1.0, 
+    kernelSizes = None, strides = None, act = "RELU"):
+    
+    layersNum = len(dims) - 1
+    
+    if kernelSizes is None:
+        kernelSizes = [config.stemKernelSize for i in range(layersNum)]
+    
+    if strides is None:
+        strides = [1 for i in range(layersNum)]
+    
+    for i in range(layersNum):
+        features = cnn(features, dims[i], dims[i+1], name = "cnn_%d" % i, batchNorm = batchNorm, 
+            dropout = dropout, kernelSize = kernelSizes[i], stride = strides[i], act = act)
+
+    return features
+
+######################################## location ########################################
+
+'''
+Computes linear positional encoding for h x w grid. 
+If outDim positive, casts positions to that dimension.
+'''
+# ignores dim
+# h,w can be tensor scalars
+def locationL(h, w, dim, outDim = -1, addBias = True):
+    dim = 2
+    grid = tf.stack(tf.meshgrid(tf.linspace(-config.locationBias, config.locationBias, w), 
+                                tf.linspace(-config.locationBias, config.locationBias, h)), axis = -1)
+
+    if outDim > 0:
+        grid = linear(grid, dim, outDim, addBias = addBias, name = "locationL")
+        dim = outDim
+
+    return grid, dim
+
+'''
+Computes sin/cos positional encoding for h x w x (4*dim). 
+If outDim positive, casts positions to that dimension.
+Based on positional encoding presented in "Attention is all you need"
+'''
+# dim % 4 = 0
+# h,w can be tensor scalars
+def locationPE(h, w, dim, outDim = -1, addBias = True):    
+    x = tf.expand_dims(tf.to_float(tf.linspace(-config.locationBias, config.locationBias, w)), axis = -1)
+    y = tf.expand_dims(tf.to_float(tf.linspace(-config.locationBias, config.locationBias, h)), axis = -1)
+    i = tf.expand_dims(tf.to_float(tf.range(dim)), axis = 0)
+
+    peSinX = tf.sin(x / (tf.pow(10000.0, i / dim)))
+    peCosX = tf.cos(x / (tf.pow(10000.0, i / dim)))
+    peSinY = tf.sin(y / (tf.pow(10000.0, i / dim)))
+    peCosY = tf.cos(y / (tf.pow(10000.0, i / dim)))
+
+    peSinX = tf.tile(tf.expand_dims(peSinX, axis = 0), [h, 1, 1])
+    peCosX = tf.tile(tf.expand_dims(peCosX, axis = 0), [h, 1, 1])
+    peSinY = tf.tile(tf.expand_dims(peSinY, axis = 1), [1, w, 1])
+    peCosY = tf.tile(tf.expand_dims(peCosY, axis = 1), [1, w, 1]) 
+
+    grid = tf.concat([peSinX, peCosX, peSinY, peCosY], axis = -1)
+    dim *= 4
+    
+    if outDim > 0:
+        grid = linear(grid, dim, outDim, addBias = addBias, name = "locationPE")
+        dim = outDim
+
+    return grid, dim
+
+locations = {
+    "L": locationL,
+    "PE": locationPE
+}
+
+'''
+Adds positional encoding to features. May ease spatial reasoning.
+(although not used in the default model). 
+
+Args:
+    features: features to add position encoding to.
+    [batchSize, h, w, c]
+
+    inDim: number of features' channels
+    lDim: dimension for positional encodings
+    outDim: if positive, cast enhanced features (with positions) to that dimension
+    h: features' height
+    w: features' width
+    locType: L for linear encoding, PE for cos/sin based positional encoding
+    mod: way to add positional encoding: concatenation (CNCT), addition (ADD), 
+            multiplication (MUL), linear transformation (LIN).
+'''
+mods = ["CNCT", "ADD", "LIN", "MUL"]
+# if outDim = -1, then will be set based on inDim, lDim
+def addLocation(features, inDim, lDim, outDim = -1, h = None, w = None, 
+    locType = "L", mod = "CNCT", name = "", reuse = None): # h,w not needed
+    
+    with tf.variable_scope("addLocation" + name, reuse = reuse):
+        batchSize = tf.shape(features)[0]
+        if h is None:
+            h = tf.shape(features)[1]
+        if w is None:
+            w = tf.shape(features)[2]
+        dim = inDim
+
+        if mod == "LIN":
+            if outDim < 0:
+                outDim = dim
+
+            grid, _ = locations[locType](h, w, lDim, outDim = outDim, addBias = False)
+            features = linear(features, dim, outDim, name = "LIN")
+            features += grid  
+            return features, outDim
+
+        if mod == "CNCT":
+            grid, lDim = locations[locType](h, w, lDim)
+            # grid = tf.zeros_like(features) + grid
+            grid = tf.tile(tf.expand_dims(grid, axis = 0), [batchSize, 1, 1, 1])
+            features = tf.concat([features, grid], axis = -1)
+            dim += lDim
+
+        elif mod == "ADD":
+            grid, _ = locations[locType](h, w, lDim, outDim = dim)
+            features += grid    
+        
+        elif mod == "MUL": # MUL
+            grid, _ = locations[locType](h, w, lDim, outDim = dim)
+
+            if outDim < 0:
+                outDim = dim
+
+            grid = tf.tile(tf.expand_dims(grid, axis = 0), [batchSize, 1, 1, 1])
+            features = tf.concat([features, grid, features * grid], axis = -1)
+            dim *= 3                
+
+        if outDim > 0:
+            features = linear(features, dim, outDim)
+            dim = outDim 
+
+    return features, dim
+
+# config.locationAwareEnd
+# H, W, _ = config.imageDims
+# projDim = config.stemProjDim
+# k = config.stemProjPooling
+# projDim on inDim or on out
+# inDim = tf.shape(features)[3]
+
+'''
+Linearize 2d image to linear vector.
+
+Args:
+    features: batch of 2d images. 
+    [batchSize, h, w, inDim]
+
+    h: image height
+
+    w: image width
+
+    inDim: number of channels
+
+    projDim: if not None, project image to that dimension before linearization
+
+    outDim: if not None, project image to that dimension after linearization
+
+    loc: if not None, add positional encoding:
+        locType: L for linear encoding, PE for cos/sin based positional encoding
+        mod: way to add positional encoding: concatenation (CNCT), addition (ADD), 
+            multiplication (MUL), linear transformation (LIN).
+        pooling: number to pool image with before linearization.
+
+Returns linearized image:
+[batchSize, outDim] (or [batchSize, (h / pooling) * (w /pooling) * projDim] if outDim not supported) 
+'''
+# loc = {"locType": str, "mod": str}
+def linearizeFeatures(features, h, w, inDim, projDim = None, outDim = None, 
+    loc = None, pooling = None):
+    
+    if pooling is None:
+        pooling = config.imageLinPool
+    
+    if loc is not None:
+        features = addLocation(features, inDim, lDim = inDim, outDim = inDim, 
+            locType = loc["locType"], mod = loc["mod"])
+
+    if projDim is not None:
+        features = linear(features, dim, projDim)
+        features = relu(features)
+        dim = projDim
+
+    if pooling > 1:
+        poolingDims = [1, pooling, pooling, 1]
+        features = tf.nn.max_pool(features, ksize = poolingDims, strides = poolingDims, 
+            padding = "SAME")
+        h /= pooling
+        w /= pooling
+  
+    dim = h * w * dim  
+    features = tf.reshape(features, (-1, dim))
+    
+    if outDim is not None:
+        features = linear(features, dim, outDim)
+        dim = outDim
+
+    return features, dim
+
+################################### multiplication ###################################
+# specific dim / proj for x / y
+'''
+"Enhanced" hadamard product between x and y:
+1. Supports optional projection of x, and y prior to multiplication.
+2. Computes simple multiplication, or a parametrized one, using diagonal of complete matrix (bi-linear) 
+3. Optionally concatenate x or y or their projection to the multiplication result.
+
+Support broadcasting
+
+Args:
+    x: left-hand side argument
+    [batchSize, dim]
+
+    y: right-hand side argument
+    [batchSize, dim]
+
+    dim: input dimension of x and y
+    
+    dropout: dropout value to apply on x and y
+
+    proj: if not None, project x and y:
+        dim: projection dimension
+        shared: use same projection for x and y
+        dropout: dropout to apply to x and y if projected
+
+    interMod: multiplication type:
+        "MUL": x * y
+        "DIAG": x * W * y for a learned diagonal parameter W
+        "BL": x' W y for a learned matrix W
+
+    concat: if not None, concatenate x or y or their projection. 
+    
+    mulBias: optional bias to stabilize multiplication (x * bias) (y * bias)
+
+Returns the multiplication result
+[batchSize, outDim] when outDim depends on the use of proj and cocnat arguments.
+'''
+# proj = {"dim": int, "shared": bool, "dropout": float} # "act": str, "actDropout": float
+## interMod = ["direct", "scalarW", "bilinear"] # "additive"
+# interMod = ["MUL", "DIAG", "BL", "ADD"]
+# concat = {"x": bool, "y": bool, "proj": bool}
+def mul(x, y, dim, dropout = 1.0, proj = None, interMod = "MUL", concat = None, mulBias = None,
+    extendY = True, name = "", reuse = None):
+    
+    with tf.variable_scope("mul" + name, reuse = reuse):                
+        origVals = {"x": x, "y": y, "dim": dim}
+
+        x = tf.nn.dropout(x, dropout)
+        y = tf.nn.dropout(y, dropout)
+        # projection
+        if proj is not None:
+            x = tf.nn.dropout(x, proj.get("dropout", 1.0))
+            y = tf.nn.dropout(y, proj.get("dropout", 1.0))
+
+            if proj["shared"]:
+                xName, xReuse = "proj", None
+                yName, yReuse = "proj", True
+            else:
+                xName, xReuse = "projX", None
+                yName, yReuse = "projY", None
+
+            x = linear(x, dim, proj["dim"], name = xName, reuse = xReuse)
+            y = linear(y, dim, proj["dim"], name = yName, reuse = yReuse)
+            dim = proj["dim"]
+            projVals = {"x": x, "y": y, "dim": dim}
+            proj["x"], proj["y"] = x, y
+
+        if extendY:
+            y = tf.expand_dims(y, axis = -2)
+            # broadcasting to have the same shape
+            y = tf.zeros_like(x) + y
+
+        # multiplication
+        if interMod == "MUL":
+            if mulBias is None:
+                mulBias = config.mulBias             
+            output = (x + mulBias) * (y + mulBias)
+        elif interMod == "DIAG":
+            W = getWeight((dim, )) # change initialization?
+            b = getBias((dim, ))
+            activations = x * W * y + b
+        elif interMod == "BL":
+            W = getWeight((dim, dim))
+            b = getBias((dim, ))            
+            output = multiply(x, W) * y + b
+        else: # "ADD"
+            output = tf.tanh(x + y)
+        # concatenation
+        if concat is not None:
+            concatVals = projVals if concat.get("proj", False) else origVals
+            if concat.get("x", False):
+                output = tf.concat([output, concatVals["x"]], axis = -1)
+                dim += concatVals["dim"]
+
+            if concat.get("y", False):
+                output = ops.concat(output, concatVals["y"], extendY = extendY)
+                dim += concatVals["dim"]
+
+    return output, dim
+
+######################################## rnns ########################################
+
+'''
+Creates an RNN cell.
+
+Args:
+    hdim: the hidden dimension of the RNN cell.
+    
+    reuse: whether the cell should reuse parameters or create new ones.
+    
+    cellType: the cell type 
+    RNN, GRU, LSTM, MiGRU, MiLSTM, ProjLSTM
+
+    act: the cell activation
+    NON, TANH, SIGMOID, RELU, ELU
+
+    projDim: if ProjLSTM, the dimension for the states projection
+
+Returns the cell.
+'''
+# tf.nn.rnn_cell.MultiRNNCell([cell(hDim, reuse = reuse) for _ in config.encNumLayers])
+# note that config.enc params not general 
+def createCell(hDim, reuse, cellType = None, act = None, projDim = None):
+    if cellType is None:
+        cellType = config.encType
+
+    activation = activations.get(act, None) 
+
+    if cellType == "ProjLSTM":
+        cell = tf.nn.rnn_cell.LSTMCell
+        if projDim is None:
+            projDim = config.cellDim
+        cell = cell(hDim, num_proj = projDim, reuse = reuse, activation = activation)
+        return cell        
+
+    cells = {
+        "RNN": tf.nn.rnn_cell.BasicRNNCell,
+        "GRU": tf.nn.rnn_cell.GRUCell,
+        "LSTM": tf.nn.rnn_cell.BasicLSTMCell,
+        "MiGRU": MiGRUCell,
+        "MiLSTM": MiLSTMCell
+    }
+
+    cell = cells[cellType](hDim, reuse = reuse, activation = activation)
+
+    return cell
+
+'''
+Runs an forward RNN layer.
+
+Args:
+    inSeq: the input sequence to run the RNN over.
+    [batchSize, sequenceLength, inDim]
+    
+    seqL: the sequence matching lengths.
+    [batchSize, 1]
+
+    hDim: hidden dimension of the RNN.
+
+    cellType: the cell type 
+    RNN, GRU, LSTM, MiGRU, MiLSTM, ProjLSTM
+
+    dropout: value for dropout over input sequence
+
+    varDp: if not None, state and input variational dropouts to apply.
+    dimension of input has to be supported (inputSize). 
+
+Returns the outputs sequence and final RNN state.  
+'''
+# varDp = {"stateDp": float, "inputDp": float, "inputSize": int}
+# proj = {"output": bool, "state": bool, "dim": int, "dropout": float, "act": str}
+def fwRNNLayer(inSeq, seqL, hDim, cellType = None, dropout = 1.0, varDp = None, 
+    name = "", reuse = None): # proj = None
+    
+    with tf.variable_scope("rnnLayer" + name, reuse = reuse):
+        batchSize = tf.shape(inSeq)[0]
+
+        cell = createCell(hDim, reuse, cellType) # passing reuse isn't mandatory
+
+        if varDp is not None:
+            cell = tf.contrib.rnn.DropoutWrapper(cell, 
+                state_keep_prob = varDp["stateDp"],
+                input_keep_prob = varDp["inputDp"],
+                variational_recurrent = True, input_size = varDp["inputSize"], dtype = tf.float32)
+        else:           
+            inSeq = tf.nn.dropout(inSeq, dropout)
+        
+        initialState = cell.zero_state(batchSize, tf.float32)
+
+        outSeq, lastState = tf.nn.dynamic_rnn(cell, inSeq, 
+            sequence_length = seqL, 
+            initial_state = initialState,
+            swap_memory = True)
+            
+        if isinstance(lastState, tf.nn.rnn_cell.LSTMStateTuple):
+            lastState = lastState.h
+
+        # if proj is not None:
+        #     if proj["output"]:
+        #         outSeq = linear(outSeq, cell.output_size, proj["dim"], act = proj["act"],  
+        #             dropout = proj["dropout"], name = "projOutput")
+
+        #     if proj["state"]:
+        #         lastState = linear(lastState, cell.state_size, proj["dim"], act = proj["act"],  
+        #             dropout = proj["dropout"], name = "projState")
+
+    return outSeq, lastState
+
+'''
+Runs an bidirectional RNN layer.
+
+Args:
+    inSeq: the input sequence to run the RNN over.
+    [batchSize, sequenceLength, inDim]
+    
+    seqL: the sequence matching lengths.
+    [batchSize, 1]
+
+    hDim: hidden dimension of the RNN.
+
+    cellType: the cell type 
+    RNN, GRU, LSTM, MiGRU, MiLSTM
+
+    dropout: value for dropout over input sequence
+
+    varDp: if not None, state and input variational dropouts to apply.
+    dimension of input has to be supported (inputSize).   
+
+Returns the outputs sequence and final RNN state.     
+'''
+# varDp = {"stateDp": float, "inputDp": float, "inputSize": int}
+# proj = {"output": bool, "state": bool, "dim": int, "dropout": float, "act": str}
+def biRNNLayer(inSeq, seqL, hDim, cellType = None, dropout = 1.0, varDp = None, 
+    name = "", reuse = None): # proj = None, 
+
+    with tf.variable_scope("birnnLayer" + name, reuse = reuse):
+        batchSize = tf.shape(inSeq)[0]
+
+        with tf.variable_scope("fw"):
+            cellFw = createCell(hDim, reuse, cellType)
+        with tf.variable_scope("bw"):
+            cellBw = createCell(hDim, reuse, cellType)
+        
+        if varDp is not None:
+            cellFw = tf.contrib.rnn.DropoutWrapper(cellFw, 
+                state_keep_prob = varDp["stateDp"],
+                input_keep_prob = varDp["inputDp"],
+                variational_recurrent = True, input_size = varDp["inputSize"], dtype = tf.float32)
+            
+            cellBw = tf.contrib.rnn.DropoutWrapper(cellBw, 
+                state_keep_prob = varDp["stateDp"],
+                input_keep_prob = varDp["inputDp"],
+                variational_recurrent = True, input_size = varDp["inputSize"], dtype = tf.float32)            
+        else:
+            inSeq = tf.nn.dropout(inSeq, dropout)
+
+        initialStateFw = cellFw.zero_state(batchSize, tf.float32)
+        initialStateBw = cellBw.zero_state(batchSize, tf.float32)
+
+        (outSeqFw, outSeqBw), (lastStateFw, lastStateBw) = tf.nn.bidirectional_dynamic_rnn(
+            cellFw, cellBw, inSeq, 
+            sequence_length = seqL, 
+            initial_state_fw = initialStateFw, 
+            initial_state_bw = initialStateBw,
+            swap_memory = True)
+
+        if isinstance(lastStateFw, tf.nn.rnn_cell.LSTMStateTuple):
+            lastStateFw = lastStateFw.h # take c? 
+            lastStateBw = lastStateBw.h  
+
+        outSeq = tf.concat([outSeqFw, outSeqBw], axis = -1)
+        lastState = tf.concat([lastStateFw, lastStateBw], axis = -1)
+
+        # if proj is not None:
+        #     if proj["output"]:
+        #         outSeq = linear(outSeq, cellFw.output_size + cellFw.output_size, 
+        #             proj["dim"], act = proj["act"], dropout = proj["dropout"], 
+        #             name = "projOutput")
+
+        #     if proj["state"]:
+        #         lastState = linear(lastState, cellFw.state_size + cellFw.state_size, 
+        #             proj["dim"], act = proj["act"], dropout = proj["dropout"], 
+        #             name = "projState")
+
+    return outSeq, lastState
+
+# int(hDim / 2) for biRNN?
+'''
+Runs an RNN layer by calling biRNN or fwRNN.
+
+Args:
+    inSeq: the input sequence to run the RNN over.
+    [batchSize, sequenceLength, inDim]
+    
+    seqL: the sequence matching lengths.
+    [batchSize, 1]
+
+    hDim: hidden dimension of the RNN.
+
+    bi: true to run bidirectional rnn.
+
+    cellType: the cell type 
+    RNN, GRU, LSTM, MiGRU, MiLSTM
+
+    dropout: value for dropout over input sequence
+
+    varDp: if not None, state and input variational dropouts to apply.
+    dimension of input has to be supported (inputSize).   
+
+Returns the outputs sequence and final RNN state.     
+'''
+# proj = {"output": bool, "state": bool, "dim": int, "dropout": float, "act": str}
+# varDp = {"stateDp": float, "inputDp": float, "inputSize": int}
+def RNNLayer(inSeq, seqL, hDim, bi = None, cellType = None, dropout = 1.0, varDp = None, 
+    name = "", reuse = None): # proj = None
+    
+    with tf.variable_scope("rnnLayer" + name, reuse = reuse):
+        if bi is None:
+            bi = config.encBi
+        
+        rnn = biRNNLayer if bi else fwRNNLayer
+        
+        if bi:
+            hDim = int(hDim / 2)
+
+    return rnn(inSeq, seqL, hDim, cellType = cellType, dropout = dropout, varDp = varDp) # , proj = proj
+
+# tf counterpart?
+# hDim = config.moduleDim
+def multigridRNNLayer(featrues, h, w, dim, name = "", reuse = None):
+    with tf.variable_scope("multigridRNNLayer" + name, reuse = reuse):
+        featrues = linear(featrues, dim, dim / 2, name = "i")
+
+        output0 = gridRNNLayer(featrues, h, w, dim, right = True, down = True, name = "rd")
+        output1 = gridRNNLayer(featrues, h, w, dim, right = True, down = False, name = "r")
+        output2 = gridRNNLayer(featrues, h, w, dim, right = False, down = True, name = "d")
+        output3 = gridRNNLayer(featrues, h, w, dim, right = False, down = False, name = "NON")
+
+        output = tf.concat([output0, output1, output2, output3], axis = -1)
+        output = linear(output, 2 * dim, dim, name = "o")
+
+    return outputs
+
+# h,w should be constants
+def gridRNNLayer(features, h, w, dim, right, down, name = "", reuse = None):
+    with tf.variable_scope("gridRNNLayer" + name):
+        batchSize = tf.shape(features)[0]
+
+        cell = createCell(dim, reuse = reuse, cellType = config.stemGridRnnMod, 
+            act = config.stemGridAct)
+        
+        initialState = cell.zero_state(batchSize, tf.float32)
+        
+        inputs = [tf.unstack(row, w, axis = 1) for row in tf.unstack(features, h, axis = 1)]
+        states = [[None for _ in range(w)] for _ in range(h)]
+
+        iAxis = range(h) if down else (range(h)[::-1])
+        jAxis = range(w) if right else (range(w)[::-1])
+
+        iPrev = -1 if down else 1
+        jPrev = -1 if right else 1
+
+        prevState = lambda i,j: states[i][j] if (i >= 0 and i < h and j >= 0 and j < w) else initialState
+        
+        for i in iAxis:
+            for j in jAxis:
+                prevs = tf.concat((prevState(i + iPrev, j), prevState(i, j + jPrev)), axis = -1)
+                curr = inputs[i][j]
+                _, states[i][j] = cell(prevs, curr)
+
+        outputs = [tf.stack(row, axis = 1) for row in states]
+        outputs = tf.stack(outputs, axis = 1)
+
+    return outputs
+
+# tf seq2seq?
+# def projRNNLayer(inSeq, seqL, hDim, labels, labelsNum, labelsDim, labelsEmb, name = "", reuse = None):
+#     with tf.variable_scope("projRNNLayer" + name):
+#         batchSize = tf.shape(features)[0]
+
+#         cell = createCell(hDim, reuse = reuse)
+
+#         projCell = ProjWrapper(cell, labelsNum, labelsDim, labelsEmb, # config.wrdEmbDim
+#             feedPrev = True, dropout = 1.0, config,
+#             temperature = 1.0, clevr_sample = False, reuse)
+        
+#         initialState = projCell.zero_state(batchSize, tf.float32)
+        
+#         if config.soft:
+#             inSeq = inSeq
+
+#             # outputs, _ = tf.nn.static_rnn(projCell, inputs, 
+#             #     sequence_length = seqL, 
+#             #     initial_state = initialState)
+
+#             inSeq = tf.unstack(inSeq, axis = 1)                        
+#             state = initialState
+#             logitsList = []
+#             chosenList = []
+
+#             for inp in inSeq:
+#                 (logits, chosen), state = projCell(inp, state)
+#                 logitsList.append(logits)
+#                 chosenList.append(chosen)
+#                 projCell.reuse = True
+
+#             logitsOut = tf.stack(logitsList, axis = 1)
+#             chosenOut = tf.stack(chosenList, axis = 1)
+#             outputs = (logitsOut, chosenOut)
+#         else:
+#             labels = tf.to_float(labels)
+#             labels = tf.concat([tf.zeros((batchSize, 1)), labels], axis = 1)[:, :-1] # ,newaxis
+#             inSeq = tf.concat([inSeq, tf.expand_dims(labels, axis = -1)], axis = -1)
+
+#             outputs, _ = tf.nn.dynamic_rnn(projCell, inSeq, 
+#                 sequence_length = seqL, 
+#                 initial_state = initialState,
+#                 swap_memory = True)
+
+#     return outputs #, labelsEmb
+
+############################### variational dropout ###############################
+
+'''
+Generates a variational dropout mask for a given shape and a dropout 
+probability value.
+'''
+def generateVarDpMask(shape, keepProb):
+    randomTensor = tf.to_float(keepProb)
+    randomTensor += tf.random_uniform(shape, minval = 0, maxval = 1)
+    binaryTensor = tf.floor(randomTensor)
+    mask = tf.to_float(binaryTensor)
+    return mask
+
+'''
+Applies the a variational dropout over an input, given dropout mask 
+and a dropout probability value. 
+'''
+def applyVarDpMask(inp, mask, keepProb):
+    ret = (tf.div(inp, tf.to_float(keepProb))) * mask
+    return ret