# Path: nn_layers.py
# def per_word_averaging_layer(dist, dist_mask):
#     """
#     :param proj: Output of the LSTM layer
#     :param dist_mask: 4D index tensor
#     :return: The per-word averages.
#     """
# 
#     dist = dist * dist_mask
# 
#     dist = dist.dimshuffle(1, 2, 0, 3)
#     dist_mask = tensor.cast(dist_mask, theano.config.floatX)
#     dist_mask = dist_mask.dimshuffle(1, 2, 0, 3)
#     divider = dist_mask.sum(axis=0)
#     divider += tensor.eq(divider, numpy_floatX(0.0))
# 
#     dist = dist * dist_mask
#     tmp = tensor.cast(dist.sum(axis=0), theano.config.floatX)
#     tmp /= divider
#     return tmp.dimshuffle(1, 0, 2)


import unittest
import theano
import theano.tensor
import numpy
import logging
from nn_layers import per_word_averaging_layer



class WordAveragingOpTests(unittest.TestCase):
    @classmethod
    def get_std(cls):
        maxw, maxwlen = 2, 3
        # n_chars -> maximum character sequence length (e.g. 140)
        # n_samples = minibatch size
        # n_proj = number of combined word/character embeddings
        n_chars, n_samples, n_proj = 4, 2, 4
        # Allocate a matrix which represents the LSTM layer output
        L = numpy.zeros((maxw, maxwlen, n_samples, n_proj))
        L[0, 0, 0, :] = [0.2, 0.3, 0.4, 0.5]
        L[0, 1, 0, :] = [-0.2, -0.4, 0.6, 0.7]
        L[1, 0, 0, :] = [0.8, -0.2, 0.3, -0.1]

# Path: nn_layers.py
# def per_word_averaging_layer(dist, dist_mask):
#     """
#     :param proj: Output of the LSTM layer
#     :param dist_mask: 4D index tensor
#     :return: The per-word averages.
#     """
# 
#     dist = dist * dist_mask
# 
#     dist = dist.dimshuffle(1, 2, 0, 3)
#     dist_mask = tensor.cast(dist_mask, theano.config.floatX)
#     dist_mask = dist_mask.dimshuffle(1, 2, 0, 3)
#     divider = dist_mask.sum(axis=0)
#     divider += tensor.eq(divider, numpy_floatX(0.0))
# 
#     dist = dist * dist_mask
#     tmp = tensor.cast(dist.sum(axis=0), theano.config.floatX)
#     tmp /= divider
#     return tmp.dimshuffle(1, 0, 2)


# Path: test_nn_layers.py
import unittest
import theano
import theano.tensor
import numpy
import logging
from nn_layers import per_word_averaging_layer




class WordAveragingOpTests(unittest.TestCase):
    @classmethod
    def get_std(cls):
        maxw, maxwlen = 2, 3
        # n_chars -> maximum character sequence length (e.g. 140)
        # n_samples = minibatch size
        # n_proj = number of combined word/character embeddings
        n_chars, n_samples, n_proj = 4, 2, 4
        # Allocate a matrix which represents the LSTM layer output
        L = numpy.zeros((maxw, maxwlen, n_samples, n_proj))
        L[0, 0, 0, :] = [0.2, 0.3, 0.4, 0.5]
        L[0, 1, 0, :] = [-0.2, -0.4, 0.6, 0.7]
        L[1, 0, 0, :] = [0.8, -0.2, 0.3, -0.1]

        L[0, 0, 1, :] = [-0.4, 0.2, 0.8, 0.9]