Make opt work when shape has been specified.

theano/tensor/nnet/nnet.py

@@ -1378,7 +1378,21 @@ def local_argmax_pushdown(node):
 
 
 def _check_rows_is_arange_len_labels(rows, labels):
-    '''Check that 'rows' is the same node as T.arange(labels.shape[0])'''
+    '''Check that 'rows' is the same node as T.arange(labels.shape[0])
+
+    Also considers the case where labels.shape[0] is constant and equal
+    to 1, and T.arange(labels.shape[0]) has been constant-folded into 0.
+    '''
+
+    if labels.owner and hasattr(labels.owner.fgraph, 'shape_feature'):
+        shape_of = labels.owner.fgraph.shape_feature.shape_of
+        # TODO: consider cases where shape_of[labels] is constant, and
+        # has a value different from 1.
+        # This case is harder, as _is_const only accepts a scalar value
+        # as second argument, so checking for
+        # _is_const(rows, numpy.arange(...)) does not work for the moment.
+        if len(shape_of[labels]) == 1 and _is_const(shape_of[labels][0], 1):
+            return _is_const(rows, 0)
 
     if rows.owner and isinstance(rows.owner.op, tensor.ARange):
         start, stop, step = rows.owner.inputs

theano/tensor/nnet/tests/test_nnet.py

@@ -925,7 +925,69 @@ class T_CrossentropyCategorical1Hot(utt.InferShapeTester):
                 config.warn.sum_div_dimshuffle_bug = backup
 
             if verbose:
-                print_graph(g)
+                self.print_graph(g)
+            try:
+                ops = [node.op for node in g.maker.fgraph.toposort()]
+                assert len(ops) <= 6
+                assert crossentropy_softmax_1hot_with_bias_dx in ops
+                assert softmax_with_bias in ops
+                assert softmax_grad not in ops
+                g(x_val, b_val, y_val)
+            except Exception:
+                theano.printing.debugprint(g)
+                raise
+
+    def test_optimize_xent_vector4(self):
+        # Same as test_optimize_xent_vector2, but y is the result of
+        # a "specify_shape" that indicates its length is 1, so the
+        # constant-folding of arange(y.shape[0]) happen before the xent
+        # optimization
+        verbose = 0
+        mode = theano.compile.mode.get_default_mode()
+        if mode == theano.compile.mode.get_mode('FAST_COMPILE'):
+            mode = 'FAST_RUN'
+        rng = numpy.random.RandomState(utt.fetch_seed())
+        x_val = rng.randn(5).astype(config.floatX)
+        b_val = rng.randn(5).astype(config.floatX)
+        y_val = numpy.asarray([2])
+
+        x = T.vector('x')
+        b = T.vector('b')
+        y_ = T.lvector('y_')
+        y = T.specify_shape(y_, (1,))
+
+        ## Test that a biased softmax is optimized correctly
+        bias_expressions = [
+                T.sum(-T.log(softmax(x + b)[T.arange(y.shape[0]), y])),
+                -T.sum(T.log(softmax(b + x)[T.arange(y.shape[0]), y])),
+                -T.sum(T.log(softmax(x + b))[T.arange(y.shape[0]), y]),
+                T.sum(-T.log(softmax(b + x))[T.arange(y.shape[0]), y])]
+
+        for expr in bias_expressions:
+            f = theano.function([x, b, y_], expr, mode=mode)
+            if verbose:
+                self.print_graph(f)
+            try:
+                ops = [node.op for node in f.maker.fgraph.toposort()]
+                # [big_op, sum, dim_shuffle, specify_shape]
+                assert len(ops) <= 4
+                assert crossentropy_softmax_argmax_1hot_with_bias in ops
+                assert not [1 for o in ops
+                            if isinstance(o, T.AdvancedSubtensor)]
+                f(x_val, b_val, y_val)
+            except Exception:
+                theano.printing.debugprint(f)
+                raise
+
+            backup = config.warn.sum_div_dimshuffle_bug
+            config.warn.sum_div_dimshuffle_bug = False
+            try:
+                g = theano.function([x, b, y], T.grad(expr, x), mode=mode)
+            finally:
+                config.warn.sum_div_dimshuffle_bug = backup
+
+            if verbose:
+                self.print_graph(g)
             try:
                 ops = [node.op for node in g.maker.fgraph.toposort()]
                 assert len(ops) <= 6