[CRASH] Make MaxAndArgmax support NoneConst as axis.

theano/tensor/basic.py

@@ -1238,7 +1238,7 @@ class MaxAndArgmax(Op):
                     axis = range(x.type.ndim)
                 assert axis == range(x.type.ndim), (
                     "MaxAndArgmax does not support multiple"
-                    " axes. the max fct supports it.")
+                    " axes. the max fct supports it. Got %s" % axis)
                 axis = None
             else:
                 axis = axis[0]
@@ -1248,8 +1248,12 @@ class MaxAndArgmax(Op):
         elif isinstance(axis, numpy.ndarray) and axis.ndim == 0:
             axis = int(axis)
         elif isinstance(axis, Variable):
-            if not isinstance(axis, TensorConstant):
-                raise TypeError("MaxAndArgmax needs a constant axis")
+            if NoneConst.equals(axis):
+                axis = None
+            elif not isinstance(axis, TensorConstant):
+                raise TypeError(
+                    "MaxAndArgmax needs a constant axis. Got %s" % axis)
+            else:
                 assert (axis.dtype.startswith("int")
                         or axis.dtype.startswith("uint"))
                 axis = int(axis.data)

theano/tensor/tests/test_basic.py

@@ -45,7 +45,7 @@ from theano.tensor import (_shared, wvector, bvector, autocast_float_as,
         itensor3, Tile, switch, Diagonal, Diag,
         nonzero, flatnonzero, nonzero_values,
         stacklists, DimShuffle, hessian, ptp, power,
-        swapaxes, choose, Choose
+        swapaxes, choose, Choose, NoneConst,
         )
 
 from theano.tests import unittest_tools as utt
@@ -2722,6 +2722,7 @@ class T_max_and_argmax(unittest.TestCase):
         n = as_tensor_variable(data)
         for (axis, np_axis)  in [(-1, -1), (0, 0), (1, 1), (None, None),
                                  ([0, 1], None), ([1, 0], None),
+                                 (NoneConst.clone(), None),
                                  (constant(0), 0)]:
             v, i = eval_outputs(max_and_argmax(n, axis))
             assert i.dtype == 'int64'