Add softmax and log_softmax to math module

Author: ricardoV94

RELEASE-NOTES.md

@@ -125,6 +125,7 @@ This includes API changes we did not warn about since at least `3.11.0` (2021-01
   - `pm.Data` now passes additional kwargs to `aesara.shared`/`at.as_tensor`. [#5098](https://github.com/pymc-devs/pymc/pull/5098).
 - Univariate censored distributions are now available via `pm.Censored`. [#5169](https://github.com/pymc-devs/pymc/pull/5169)
 - Nested models now inherit the parent model's coordinates. [#5344](https://github.com/pymc-devs/pymc/pull/5344)
+- `softmax` and `log_softmax` functions added to `math` module (see [#5279](https://github.com/pymc-devs/pymc/pull/5279)).
 - ...
 
 

pymc/math.py

@@ -211,6 +211,22 @@ def invlogit(x, eps=None):
     return at.sigmoid(x)
 
 
+def softmax(x, axis=None):
+    # Ignore vector case UserWarning issued by Aesara. This can be removed once Aesara
+    # drops that warning
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", UserWarning)
+        return at.nnet.softmax(x, axis=axis)
+
+
+def log_softmax(x, axis=None):
+    # Ignore vector case UserWarning issued by Aesara. This can be removed once Aesara
+    # drops that warning
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", UserWarning)
+        return at.nnet.logsoftmax(x, axis=axis)
+
+
 def logbern(log_p):
     if np.isnan(log_p):
         raise FloatingPointError("log_p can't be nan.")

pymc/tests/test_math.py

@@ -30,10 +30,12 @@
     kronecker,
     log1mexp,
     log1mexp_numpy,
+    log_softmax,
     logdet,
     logdiffexp,
     logdiffexp_numpy,
     probit,
+    softmax,
 )
 from pymc.tests.helpers import SeededTest, verify_grad
 
@@ -265,3 +267,24 @@ def test_invlogit_deprecation_warning():
     assert not record
 
     assert np.isclose(res, res_zero_eps)
+
+
+@pytest.mark.parametrize(
+    "aesara_function, pymc_wrapper",
+    [
+        (at.nnet.softmax, softmax),
+        (at.nnet.logsoftmax, log_softmax),
+    ],
+)
+def test_softmax_logsoftmax_no_warnings(aesara_function, pymc_wrapper):
+    """Test that wrappers for aesara functions do not issue Warnings"""
+
+    vector = at.vector("vector")
+    with pytest.warns(None) as record:
+        aesara_function(vector)
+    warnings = {warning.category for warning in record.list}
+    assert warnings == {UserWarning, FutureWarning}
+
+    with pytest.warns(None) as record:
+        pymc_wrapper(vector)
+    assert not record