Move test helper softmax_graph to test module

aesara/tensor/nnet/__init__.py

@@ -35,7 +35,6 @@
     sigmoid_binary_crossentropy,
     softmax,
     softmax_grad_legacy,
-    softmax_graph,
     softmax_legacy,
     softmax_simplifier,
     softmax_with_bias,

aesara/tensor/nnet/basic.py

@@ -1122,10 +1122,6 @@ def local_logsoftmax_grad(fgraph, node):
         return [ret]
 
 
-def softmax_graph(c):
-    return exp(c) / exp(c).sum(axis=-1, keepdims=True)
-
-
 UNSET_AXIS = object()
 
 

tests/scan/test_basic.py

@@ -56,7 +56,7 @@
 from aesara.tensor.math import dot, mean, sigmoid
 from aesara.tensor.math import sum as aet_sum
 from aesara.tensor.math import tanh
-from aesara.tensor.nnet import categorical_crossentropy, softmax_graph
+from aesara.tensor.nnet import categorical_crossentropy
 from aesara.tensor.random.utils import RandomStream
 from aesara.tensor.shape import Shape_i, reshape, shape, specify_shape
 from aesara.tensor.sharedvar import SharedVariable
@@ -81,6 +81,7 @@
     vector,
 )
 from tests import unittest_tools as utt
+from tests.tensor.nnet.test_basic import softmax_graph
 
 
 if config.mode == "FAST_COMPILE":

tests/tensor/nnet/test_basic.py

@@ -52,7 +52,6 @@
     sigmoid_binary_crossentropy,
     softmax,
     softmax_grad_legacy,
-    softmax_graph,
     softmax_legacy,
     softmax_with_bias,
     softsign,
@@ -83,6 +82,10 @@
 )
 
 
+def softmax_graph(c):
+    return exp(c) / exp(c).sum(axis=-1, keepdims=True)
+
+
 def valid_axis_tester(Op):
     with pytest.raises(TypeError):
         Op(1.5)