test softmax_op and softmax_ghraph in test_argmax_pushdown

theano/tensor/nnet/tests/test_nnet.py

@@ -1127,10 +1127,10 @@ class T_CrossentropyCategorical1Hot(utt.InferShapeTester):
 
 def test_argmax_pushdown():
     x = tensor.matrix()
-
+    for softmax in [softmax_graph, softmax_op]:
         # test that the max_and_argmax is pushed down if the max is not used
         out = tensor.max_and_argmax(
-            softmax_graph(tensor.exp(tensor.tanh(sigmoid(x)))),
+                softmax(tensor.exp(tensor.tanh(sigmoid(x)))),
                 axis=-1)[1]
         fgraph = gof.FunctionGraph(
                 [x],
@@ -1147,7 +1147,7 @@ def test_argmax_pushdown():
         x = tensor.matrix()
         # test that the max_and_argmax is not pushed down if the max is used
         out = tensor.max_and_argmax(
-            softmax_op(tensor.exp(tensor.tanh(sigmoid(x)))),
+                softmax(tensor.exp(tensor.tanh(sigmoid(x)))),
                 axis=-1)[0]
         fgraph = gof.FunctionGraph(
                 [x],