Implement a get_constant_idx helper method on Subtensor and use it when we need…

theano/tensor/basic.py

@@ -591,20 +591,7 @@ def get_scalar_constant_value(v):
             return ret[0][0]
         if isinstance(v.owner.op, theano.tensor.subtensor.Subtensor) and v.ndim == 0:
             if isinstance(v.owner.inputs[0], TensorConstant):
-                indices = list(reversed(list(v.owner.inputs[1:])))
-                def conv(e):
-                    if isinstance(e, gof.Type):
-                        return get_scalar_constant_value(indices.pop())
-                    elif isinstance(e, slice):
-                        return slice(conv(e.start),
-                                     conv(e.stop),
-                                     conv(e.step))
-                    elif isinstance(e, (int, long, numpy.integer)):
-                        return int(e)
-                    else:
-                        raise NotScalarConstantError(v)
-                cdata = tuple(map(conv, v.owner.op.idx_list))
-
+                cdata = v.owner.op.get_constant_idx(v.owner.inputs)
                 try:
                     return v.owner.inputs[0].data.__getitem__(cdata)
                 except IndexError:

theano/tensor/nnet/nnet.py

@@ -1396,8 +1396,9 @@ def _check_rows_is_arange_len_labels(rows, labels):
         # ShapeOptimizer, but we keep it if ShapeOptimizer is not present
         if isinstance(stop.owner.op, subtensor.Subtensor):
             shape_subtensor = stop.owner
-            if list(shape_subtensor.op.idx_list) == [0]:
-                shape_var, = shape_subtensor.inputs
+            if shape_subtensor.op.get_constant_idx(shape_subtensor.inputs,
+                                                   allow_partial=True) == [0]:
+                shape_var = shape_subtensor.inputs[0]
                 if shape_var.owner and shape_var.owner.op == tensor.shape:
                     return shape_var.owner.inputs[0] is labels
         else:

theano/tensor/opt.py

@@ -1625,32 +1625,17 @@ def local_useless_subtensor(node):
         if not hasattr(node.fgraph, 'shape_feature'):
             return
         shape_of = node.fgraph.shape_feature.shape_of
-        idx_vals = get_idx_list(node.inputs, node.op.idx_list)
-        for pos, idx in enumerate(idx_vals):
+        cdata = node.op.get_constant_idx(node.inputs, allow_partial=True)
+        for pos, idx in enumerate(cdata):
             if not isinstance(idx, slice):
                 # If idx is not a slice, this means we remove this dimension
                 # from the output, so the subtensor is not useless
                 return False
-
-            # Grab the values for start/stop/step
-            try:
-                start = get_scalar_constant_value(idx.start)
-            except NotScalarConstantError:
-                start = idx.start
-            try:
-                stop = get_scalar_constant_value(idx.stop)
-            except NotScalarConstantError:
-                stop = idx.stop
-            try:
-                step = get_scalar_constant_value(idx.step)
-            except NotScalarConstantError:
-                step = idx.step
-
-            if start not in [0, None]:
+            if idx.start not in [0, None]:
                 # If the start of the slice is different from 0, or is a
                 # variable, then we assume the subtensor is not useless
                 return False
-            if step not in [1, None]:
+            if idx.step not in [1, None]:
                 # If we are going backwards, or skipping elements, then this
                 # is not a useless subtensor
                 return False
@@ -1663,11 +1648,11 @@ def local_useless_subtensor(node):
             except NotScalarConstantError:
                 pass
 
-            if isinstance(stop, (int, numpy.integer)):
+            if isinstance(idx.stop, (int, numpy.integer)):
                 if stop < length_pos_data:
                     return False
-            elif isinstance(stop, gof.Variable):
-                length_pos_shape_i = stop
+            elif isinstance(idx.stop, gof.Variable):
+                length_pos_shape_i = idx.stop
                 # length_pos is a tensor variable, but length_pos_shape_i
                 # is a scalar variable. We try to see if they represent
                 # the same underlying variable.
@@ -1694,7 +1679,7 @@ def local_useless_subtensor(node):
                 # length_pos_shape_i cannot be None
                 if length_pos_shape_i != length_pos:
                     return False
-            elif stop is None:
+            elif idx.stop is None:
                 pass
             else:
                 return False

theano/tensor/subtensor.py

@@ -362,6 +362,37 @@ class Subtensor(Op):
         else:
             raise AdvancedIndexingError(Subtensor.e_indextype, entry)
 
+    def get_constant_idx(self, inputs, allow_partial=False):
+        """
+        Return the idx_list with constant inputs replaced by their
+        python scalar equivalent.  May raise
+        `theano.tensor.NotScalarConstantError` if the idx contains
+        non-constant entries.
+
+        If allow_partial is True, then entries that are not constant
+        will stay as their input variable rather than raising an
+        exception.
+
+        None entries are always left as-is.
+        """
+        real_idx = get_idx_list(inputs, self.idx_list)
+        def conv(val):
+            if val is None:
+                return None
+            elif isinstance(val, slice):
+                return slice(conv(val.start),
+                             conv(val.stop),
+                             conv(val.step))
+            else:
+                try:
+                    return get_scalar_constant_value(val)
+                except theano.tensor.NotScalarConstantError:
+                    if allow_partial:
+                        return val
+                    else:
+                        raise
+        return map(conv, real_idx)
+
     def __init__(self, idx_list):
         self.idx_list = tuple(map(self.convert, idx_list))
         self.perform_cache_cdata = None
@@ -404,29 +435,18 @@ class Subtensor(Op):
                     % (input.type, expected_type))
 
         # infer the broadcasting pattern
-        padded = (idx_list
-                + [slice(None, None, None)] * (x.type.ndim - len(idx_list)))
-        idx_padded = get_idx_list((None,)+inputs, padded)
+        padded = (self.get_constant_idx((None,)+inputs, allow_partial=True)
+                  + [slice(None, None, None)] * (x.type.ndim - len(idx_list)))
         broadcastable = []
-        for i, (p, bc) in enumerate(izip(idx_padded, x.type.broadcastable)):
+        for i, (p, bc) in enumerate(izip(padded, x.type.broadcastable)):
             if isinstance(p, slice):
-                # figure out the value of start and stop (if they are constant)
-                try:
-                    if p.start is None:
+                if bc and p.start in [None, 0]:
+                    start = p.start
+                    if start is None:
                         start = 0
-                    else:
-                        start = get_scalar_constant_value(p.start)
-                    if p.stop is None:
-                        stop = 1
-                    else:
-                        stop = get_scalar_constant_value(p.stop)
-                except theano.tensor.NotScalarConstantError:
-                    start = None
-                    stop = None
-
-                if bc and start == 0 and stop > start:
-                    broadcastable.append(True)
-                    continue
+                    if p.stop is None or p.stop > start:
+                        broadcastable.append(True)
+                        continue
 
                 broadcastable.append(False)