Remove use of aesara.tensor.nnet in other tests

Author: rlouf

tests/link/jax/test_elemwise.py

@@ -5,10 +5,9 @@
 from aesara.graph.fg import FunctionGraph
 from aesara.graph.op import get_test_value
 from aesara.tensor import elemwise as at_elemwise
-from aesara.tensor import nnet as at_nnet
 from aesara.tensor.math import SoftmaxGrad
 from aesara.tensor.math import all as at_all
-from aesara.tensor.math import prod
+from aesara.tensor.math import log_softmax, prod, softmax
 from aesara.tensor.math import sum as at_sum
 from aesara.tensor.type import matrix, tensor, vector
 from tests.link.jax.test_basic import compare_jax_and_py
@@ -76,7 +75,7 @@ def test_jax_CAReduce():
 def test_softmax(axis):
     x = matrix("x")
     x.tag.test_value = np.arange(6, dtype=config.floatX).reshape(2, 3)
-    out = at_nnet.softmax(x, axis=axis)
+    out = softmax(x, axis=axis)
     fgraph = FunctionGraph([x], [out])
     compare_jax_and_py(fgraph, [get_test_value(i) for i in fgraph.inputs])
 
@@ -85,7 +84,7 @@ def test_softmax(axis):
 def test_logsoftmax(axis):
     x = matrix("x")
     x.tag.test_value = np.arange(6, dtype=config.floatX).reshape(2, 3)
-    out = at_nnet.logsoftmax(x, axis=axis)
+    out = log_softmax(x, axis=axis)
     fgraph = FunctionGraph([x], [out])
     compare_jax_and_py(fgraph, [get_test_value(i) for i in fgraph.inputs])
 

tests/link/jax/test_scalar.py

@@ -7,7 +7,6 @@
 from aesara.graph.fg import FunctionGraph
 from aesara.graph.op import get_test_value
 from aesara.scalar.basic import Composite
-from aesara.tensor import nnet as at_nnet
 from aesara.tensor.elemwise import Elemwise
 from aesara.tensor.math import all as at_all
 from aesara.tensor.math import (
@@ -128,10 +127,6 @@ def test_nnet():
     fgraph = FunctionGraph([x], [out])
     compare_jax_and_py(fgraph, [get_test_value(i) for i in fgraph.inputs])
 
-    out = at_nnet.ultra_fast_sigmoid(x)
-    fgraph = FunctionGraph([x], [out])
-    compare_jax_and_py(fgraph, [get_test_value(i) for i in fgraph.inputs])
-
     out = softplus(x)
     fgraph = FunctionGraph([x], [out])
     compare_jax_and_py(fgraph, [get_test_value(i) for i in fgraph.inputs])

tests/scalar/test_basic.py

@@ -444,7 +444,7 @@ def test_grad_inrange():
 
 def test_grad_abs():
     a = fscalar("a")
-    b = aesara.tensor.nnet.relu(a)
+    b = 0.5 * (a + aesara.tensor.abs(a))
     c = aesara.grad(b, a)
     f = aesara.function([a], c, mode=Mode(optimizer=None))
     # Currently Aesara return 0.5, but it isn't sure it won't change

tests/scan/test_basic.py

@@ -43,7 +43,6 @@
 from aesara.tensor.math import dot, exp, mean, sigmoid
 from aesara.tensor.math import sum as at_sum
 from aesara.tensor.math import tanh
-from aesara.tensor.nnet import categorical_crossentropy
 from aesara.tensor.random import normal
 from aesara.tensor.random.utils import RandomStream
 from aesara.tensor.shape import Shape_i, reshape, specify_shape
@@ -58,7 +57,6 @@
     fscalar,
     ftensor3,
     fvector,
-    imatrix,
     iscalar,
     ivector,
     lscalar,
@@ -3810,36 +3808,6 @@ def f_rnn_cmpl(u1_t, u2_tm1, u2_t, u2_tp1, x_tm1, y_tm1, y_tm3, W_in1):
 
         # TODO FIXME: What is this testing?  At least assert something.
 
-    def test_grad_two_scans(self):
-
-        # data input & output
-        x = tensor3("x")
-        t = imatrix("t")
-
-        # forward pass
-        W = shared(
-            np.random.default_rng(utt.fetch_seed()).random((2, 2)).astype("float32"),
-            name="W",
-            borrow=True,
-        )
-
-        def forward_scanner(x_t):
-            a2_t = dot(x_t, W)
-            y_t = softmax_graph(a2_t)
-            return y_t
-
-        y, _ = scan(fn=forward_scanner, sequences=x, outputs_info=[None])
-
-        # loss function
-        def error_scanner(y_t, t_t):
-            return mean(categorical_crossentropy(y_t, t_t))
-
-        L, _ = scan(fn=error_scanner, sequences=[y, t], outputs_info=[None])
-        L = mean(L)
-
-        # backward pass
-        grad(L, [W])
-
     def _grad_mout_helper(self, n_iters, mode):
         rng = np.random.default_rng(utt.fetch_seed())
         n_hid = 3