Merge pull request #959 from rizar/lookup_table_update

Lookup table changes

blocks/bricks/lookup.py

@@ -1,6 +1,7 @@
 """Introduces Lookup brick."""
 from blocks.bricks import Initializable
 from blocks.bricks.base import application, lazy
+from blocks.roles import WEIGHT, add_role
 from blocks.utils import check_theano_variable, shared_floatx_nans
 
 
@@ -35,11 +36,12 @@ def W(self):
     def _allocate(self):
         self.parameters.append(shared_floatx_nans((self.length, self.dim),
                                name='W'))
+        add_role(self.parameters[-1], WEIGHT)
 
     def _initialize(self):
         self.weights_init.initialize(self.W, self.rng)
 
-    @application
+    @application(inputs=['indices'], outputs=['output'])
     def apply(self, indices):
         """Perform lookup.
 
@@ -61,3 +63,9 @@ def apply(self, indices):
         output_shape = [indices.shape[i]
                         for i in range(indices.ndim)] + [self.dim]
         return self.W[indices.flatten()].reshape(output_shape)
+
+    def get_dim(self, name):
+        if name == 'output':
+            return self.dim
+        if name == 'indices':
+            return 0

tests/bricks/test_lookup.py

@@ -21,3 +21,7 @@ def test_lookup_table():
     desired = numpy.array([[[3, 4, 5], [6, 7, 8]], [[0, 1, 2], [9, 10, 11]]],
                           dtype=theano.config.floatX)
     assert_equal(f(x_val)[0], desired)
+
+    # Test get_dim
+    assert_equal(lt.get_dim(lt.apply.inputs[0]), 0)
+    assert_equal(lt.get_dim(lt.apply.outputs[0]), lt.dim)