From c674cbee48e72fdd49023a5898fc3542f3db689a Mon Sep 17 00:00:00 2001 From: Abhinav-Khot <23110006@iitgn.ac.in> Date: Thu, 26 Jun 2025 23:45:18 +0530 Subject: [PATCH 1/5] add negative value support for the digamma function --- pytensor/scalar/math.py | 23 ++++++++++++++++++++--- tests/scalar/test_math.py | 27 +++++++++++++++++++++++++++ 2 files changed, 47 insertions(+), 3 deletions(-) diff --git a/pytensor/scalar/math.py b/pytensor/scalar/math.py index 86029e626f..50ae7c6c70 100644 --- a/pytensor/scalar/math.py +++ b/pytensor/scalar/math.py @@ -389,6 +389,10 @@ def c_support_code(self, **kwargs): #define ga_double double #endif + #ifndef M_PI + #define M_PI 3.14159265358979323846 + #endif + #ifndef _PSIFUNCDEFINED #define _PSIFUNCDEFINED DEVICE double _psi(ga_double x) { @@ -396,7 +400,8 @@ def c_support_code(self, **kwargs): /*taken from Bernardo, J. M. (1976). Algorithm AS 103: Psi (Digamma) Function. Applied Statistics. 25 (3), 315-317. - http://www.uv.es/~bernardo/1976AppStatist.pdf */ + http://www.uv.es/~bernardo/1976AppStatist.pdf + */ ga_double y, R, psi_ = 0; ga_double S = 1.0e-5; @@ -406,10 +411,22 @@ def c_support_code(self, **kwargs): ga_double S5 = 3.968253968e-3; ga_double D1 = -0.5772156649; + if (x <= 0) { + // the digamma function approaches infinity from one side and -infinity from the other, around negative integers and zero + if (x == floor(x)) { + return INFINITY; // note that scipy returns -INF for 0 and NaN for negative integers + } + + // Use reflection formula + ga_double pi_x = M_PI * x; + ga_double cot_pi_x = cos(pi_x) / sin(pi_x); + return _psi(1.0 - x) + M_PI * cot_pi_x; + } + y = x; - if (y <= 0.0) - return psi_; + if (y <= 0) + return psi_; if (y <= S) return D1 - 1.0/y; diff --git a/tests/scalar/test_math.py b/tests/scalar/test_math.py index f4a9f2d414..cd008b0d6b 100644 --- a/tests/scalar/test_math.py +++ b/tests/scalar/test_math.py @@ -2,6 +2,7 @@ import numpy as np import pytest +import scipy import scipy.special as sp import pytensor.tensor as pt @@ -19,6 +20,7 @@ gammal, gammau, hyp2f1, + psi, ) from tests.link.test_link import make_function @@ -149,3 +151,28 @@ def test_scalarloop_grad_mixed_dtypes(op, scalar_loop_grads): (var.owner and isinstance(var.owner.op, ScalarLoop)) for var in ancestors(grad) ) + + +@pytest.mark.parametrize( + "linker", + ["py", "c"], +) +def test_psi(linker): + x = float64("x") + out = psi(x) + + fn = function([x], out, mode=Mode(linker=linker, optimizer="fast_run")) + fn.dprint() + + x_test = np.float64(0.5) + + np.testing.assert_allclose( + fn(x_test), + scipy.special.psi(x_test), + strict=True, + ) + np.testing.assert_allclose( + fn(-x_test), + scipy.special.psi(-x_test), + strict=True, + ) From 055d8d118ba2540a383930b4ab152777eacd1833 Mon Sep 17 00:00:00 2001 From: Abhinav-Khot <23110006@iitgn.ac.in> Date: Sat, 5 Jul 2025 01:23:45 +0530 Subject: [PATCH 2/5] Add support for negative values in psi --- pytensor/scalar/math.py | 88 ++++++++++++++++----------------------- tests/scalar/test_math.py | 4 +- 2 files changed, 39 insertions(+), 53 deletions(-) diff --git a/pytensor/scalar/math.py b/pytensor/scalar/math.py index 50ae7c6c70..e18d11dea3 100644 --- a/pytensor/scalar/math.py +++ b/pytensor/scalar/math.py @@ -378,17 +378,6 @@ def L_op(self, inputs, outputs, grads): def c_support_code(self, **kwargs): return """ - // For GPU support - #ifdef WITHIN_KERNEL - #define DEVICE WITHIN_KERNEL - #else - #define DEVICE - #endif - - #ifndef ga_double - #define ga_double double - #endif - #ifndef M_PI #define M_PI 3.14159265358979323846 #endif @@ -397,51 +386,48 @@ def c_support_code(self, **kwargs): #define _PSIFUNCDEFINED DEVICE double _psi(ga_double x) { - /*taken from - Bernardo, J. M. (1976). Algorithm AS 103: - Psi (Digamma) Function. Applied Statistics. 25 (3), 315-317. - http://www.uv.es/~bernardo/1976AppStatist.pdf - */ - - ga_double y, R, psi_ = 0; - ga_double S = 1.0e-5; - ga_double C = 8.5; - ga_double S3 = 8.333333333e-2; - ga_double S4 = 8.333333333e-3; - ga_double S5 = 3.968253968e-3; - ga_double D1 = -0.5772156649; - - if (x <= 0) { - // the digamma function approaches infinity from one side and -infinity from the other, around negative integers and zero - if (x == floor(x)) { - return INFINITY; // note that scipy returns -INF for 0 and NaN for negative integers - } + /*taken from + Bernardo, J. M. (1976). Algorithm AS 103: + Psi (Digamma) Function. Applied Statistics. 25 (3), 315-317. + http://www.uv.es/~bernardo/1976AppStatist.pdf + */ - // Use reflection formula - ga_double pi_x = M_PI * x; - ga_double cot_pi_x = cos(pi_x) / sin(pi_x); - return _psi(1.0 - x) + M_PI * cot_pi_x; - } + double y, R, psi_ = 0; + double S = 1.0e-5; + double C = 8.5; + double S3 = 8.333333333e-2; + double S4 = 8.333333333e-3; + double S5 = 3.968253968e-3; + double D1 = -0.5772156649; - y = x; + if (x <= 0) { + // the digamma function approaches infinity from one side and -infinity from the other, around negative integers and zero + if (x == floor(x)) { + return INFINITY; // note that scipy returns -INF for 0 and NaN for negative integers + } + + // Use reflection formula + ga_double pi_x = M_PI * x; + ga_double cot_pi_x = cos(pi_x) / sin(pi_x); + return _psi(1.0 - x) - M_PI * cot_pi_x; + } - if (y <= 0) - return psi_; + y = x; - if (y <= S) - return D1 - 1.0/y; + if (y <= S) + return D1 - 1.0/y; - while (y < C) { - psi_ = psi_ - 1.0 / y; - y = y + 1; - } + while (y < C) { + psi_ = psi_ - 1.0 / y; + y = y + 1; + } - R = 1.0 / y; - psi_ = psi_ + log(y) - .5 * R ; - R= R*R; - psi_ = psi_ - R * (S3 - R * (S4 - R * S5)); + R = 1.0 / y; + psi_ = psi_ + log(y) - .5 * R ; + R= R*R; + psi_ = psi_ - R * (S3 - R * (S4 - R * S5)); - return psi_; + return psi_; } #endif """ @@ -450,8 +436,8 @@ def c_code(self, node, name, inp, out, sub): (x,) = inp (z,) = out if node.inputs[0].type in float_types: - return f"""{z} = - _psi({x});""" + dtype = "npy_" + node.outputs[0].dtype + return f"{z} = ({dtype}) _psi({x});" raise NotImplementedError("only floating point is implemented") diff --git a/tests/scalar/test_math.py b/tests/scalar/test_math.py index cd008b0d6b..c4a3218ea6 100644 --- a/tests/scalar/test_math.py +++ b/tests/scalar/test_math.py @@ -155,7 +155,7 @@ def test_scalarloop_grad_mixed_dtypes(op, scalar_loop_grads): @pytest.mark.parametrize( "linker", - ["py", "c"], + ["py", "cvm"], ) def test_psi(linker): x = float64("x") @@ -164,7 +164,7 @@ def test_psi(linker): fn = function([x], out, mode=Mode(linker=linker, optimizer="fast_run")) fn.dprint() - x_test = np.float64(0.5) + x_test = np.float64(0.7) np.testing.assert_allclose( fn(x_test), From bf3016cfbd00b7b45b398218563f1f8f1d1b5970 Mon Sep 17 00:00:00 2001 From: Abhinav-Khot <23110006@iitgn.ac.in> Date: Sat, 5 Jul 2025 01:29:45 +0530 Subject: [PATCH 3/5] Add support for negative values in psi --- pytensor/scalar/math.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/pytensor/scalar/math.py b/pytensor/scalar/math.py index e18d11dea3..5dc14a5892 100644 --- a/pytensor/scalar/math.py +++ b/pytensor/scalar/math.py @@ -378,6 +378,13 @@ def L_op(self, inputs, outputs, grads): def c_support_code(self, **kwargs): return """ + // For GPU support + #ifdef WITHIN_KERNEL + #define DEVICE WITHIN_KERNEL + #else + #define DEVICE + #endif + #ifndef M_PI #define M_PI 3.14159265358979323846 #endif From 526b6c67125757f655a7c1bfc079add7550daf09 Mon Sep 17 00:00:00 2001 From: Abhinav-Khot <23110006@iitgn.ac.in> Date: Sat, 5 Jul 2025 01:41:37 +0530 Subject: [PATCH 4/5] fix type issue --- pytensor/scalar/math.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pytensor/scalar/math.py b/pytensor/scalar/math.py index 5dc14a5892..d08759a978 100644 --- a/pytensor/scalar/math.py +++ b/pytensor/scalar/math.py @@ -391,7 +391,7 @@ def c_support_code(self, **kwargs): #ifndef _PSIFUNCDEFINED #define _PSIFUNCDEFINED - DEVICE double _psi(ga_double x) { + DEVICE double _psi(double x) { /*taken from Bernardo, J. M. (1976). Algorithm AS 103: @@ -414,8 +414,8 @@ def c_support_code(self, **kwargs): } // Use reflection formula - ga_double pi_x = M_PI * x; - ga_double cot_pi_x = cos(pi_x) / sin(pi_x); + double pi_x = M_PI * x; + double cot_pi_x = cos(pi_x) / sin(pi_x); return _psi(1.0 - x) - M_PI * cot_pi_x; } From 90714db77a7d8aea3eaea34a75815b86cf8df67e Mon Sep 17 00:00:00 2001 From: Abhinav-Khot <23110006@iitgn.ac.in> Date: Sat, 5 Jul 2025 01:52:44 +0530 Subject: [PATCH 5/5] fix failing tests\ --- tests/scalar/test_math.py | 12 ++---------- 1 file changed, 2 insertions(+), 10 deletions(-) diff --git a/tests/scalar/test_math.py b/tests/scalar/test_math.py index c4a3218ea6..da116ab887 100644 --- a/tests/scalar/test_math.py +++ b/tests/scalar/test_math.py @@ -166,13 +166,5 @@ def test_psi(linker): x_test = np.float64(0.7) - np.testing.assert_allclose( - fn(x_test), - scipy.special.psi(x_test), - strict=True, - ) - np.testing.assert_allclose( - fn(-x_test), - scipy.special.psi(-x_test), - strict=True, - ) + np.testing.assert_allclose(fn(x_test), scipy.special.psi(x_test)) + np.testing.assert_allclose(fn(-x_test), scipy.special.psi(-x_test))