diff --git a/pytensor/tensor/math.py b/pytensor/tensor/math.py index 2ba2ab1587..714f597b32 100644 --- a/pytensor/tensor/math.py +++ b/pytensor/tensor/math.py @@ -2453,6 +2453,11 @@ def kv(v, x): return kve(v, x) * exp(-x) +def kn(n, x): + """Modified Bessel function of the second kind of integer order v.""" + return kv(n, x) + + @scalar_elemwise def sigmoid(x): """Logistic sigmoid function (1 / (1 + exp(-x)), also known as expit or inverse logit""" @@ -4337,6 +4342,7 @@ def nan_to_num(x, nan=0.0, posinf=None, neginf=None): "i1", "iv", "ive", + "kn", "kv", "kve", "sigmoid", diff --git a/tests/tensor/test_math_scipy.py b/tests/tensor/test_math_scipy.py index 8f70950206..e15293e979 100644 --- a/tests/tensor/test_math_scipy.py +++ b/tests/tensor/test_math_scipy.py @@ -18,7 +18,7 @@ from pytensor import tensor as pt from pytensor.compile.mode import get_default_mode from pytensor.configdefaults import config -from pytensor.tensor import gammaincc, inplace, kv, kve, vector +from pytensor.tensor import gammaincc, inplace, kn, kv, kve, vector from tests import unittest_tools as utt from tests.tensor.utils import ( _good_broadcast_unary_chi2sf, @@ -1220,3 +1220,17 @@ def test_kv(): out.eval({v: test_v, x: test_x}), scipy.special.kv(test_v[:, None], test_x[None, :]), ) + + +def test_kn(): + n = vector("n") + x = vector("x") + + out = kn(n[:, None], x[None, :]) + test_n = np.array([-3, 4, 0, 5], dtype=n.type.dtype) + test_x = np.linspace(0, 512, 10, dtype=x.type.dtype) + + np.testing.assert_allclose( + out.eval({n: test_n, x: test_x}), + scipy.special.kn(test_n[:, None], test_x[None, :]), + )