#!/usr/bin/env python
# -*- coding: utf-8 -*-

import numpy as np
import theano.tensor as T
from deepy.layers import NeuralLayer
from deepy.utils import FLOATX

class WordEmbedding(NeuralLayer):
    """
    Word embedding layer.
    The word embeddings are randomly initialized, and are learned over the time.
    """
    def __init__(self, size, vocab_size, zero_index=None):
        super(WordEmbedding, self).__init__("word_embed")
        self.size = size
        self.vocab_size = vocab_size
        self.output_dim = size
        self.zero_index = zero_index

    def setup(self):
        self.embed_matrix = self.create_weight(self.vocab_size, self.size, "embed")
        self.register_parameters(self.embed_matrix)

    def output(self, x):
        if self.zero_index != None:
            mask = T.neq(x, self.zero_index)
            # To avoid negative index
            x *= mask

        ret_tensor = self.embed_matrix[x.flatten()].reshape((x.shape[0], x.shape[1], self.size))

        if self.zero_index != None:
            ret_tensor *= mask[:, :, None]
        return ret_tensor

@@ -948,16 +948,20 @@
 ape(
+list
 (x.shape
 %5B0%5D,
@@ -960,25 +960,13 @@
 hape
-%5B0%5D, x.shape%5B1%5D, 
+) + %5B
 self
@@ -970,17 +970,17 @@
 elf.size
-)
+%5D
 )%0A%0A     
@@ -1014,32 +1014,141 @@
 ne:%0A            
+if x.ndim == 2:%0A                ret_tensor *= mask%5B:, :, None%5D%0A            elif x.ndim == 1:%0A                
 ret_tensor *= ma
@@ -1142,35 +1142,32 @@
 _tensor *= mask%5B
-:, 
 :, None%5D%0A       

Enhance word_embed