Move Softmax, LogSoftmax, SoftmaxGrad to new aesara.tensor.special

Author: rlouf

aesara/link/jax/dispatch/elemwise.py

@@ -3,7 +3,7 @@
 
 from aesara.link.jax.dispatch.basic import jax_funcify, jnp_safe_copy
 from aesara.tensor.elemwise import CAReduce, DimShuffle, Elemwise
-from aesara.tensor.math import LogSoftmax, Softmax, SoftmaxGrad
+from aesara.tensor.special import LogSoftmax, Softmax, SoftmaxGrad
 
 
 @jax_funcify.register(Elemwise)

aesara/link/numba/dispatch/elemwise.py

@@ -38,13 +38,8 @@
 from aesara.scalar.basic import add as add_as
 from aesara.scalar.basic import scalar_maximum
 from aesara.tensor.elemwise import CAReduce, DimShuffle, Elemwise
-from aesara.tensor.math import (
-    LogSoftmax,
-    MaxAndArgmax,
-    MulWithoutZeros,
-    Softmax,
-    SoftmaxGrad,
-)
+from aesara.tensor.math import MaxAndArgmax, MulWithoutZeros
+from aesara.tensor.special import LogSoftmax, Softmax, SoftmaxGrad
 
 
 @singledispatch

aesara/tensor/__init__.py

@@ -113,6 +113,7 @@ def _get_vector_length_Constant(op: Union[Op, Variable], var: Constant) -> int:
 
 # isort: off
 from aesara.tensor import linalg  # noqa
+from aesara.tensor import special
 
 # For backward compatibility
 from aesara.tensor import nlinalg  # noqa