torch.polygamma inconsistent with scipy.special.polygamma for n >= 1

[igm503]

kshitij12345 raised an issue two years ago concerning the fact that pytorch's implementation of the 1st order polygamma function incorrectly produces finite values for negative integers. The issue was closed, but the issue still seems to be present in the latest nightlies and the latest release of pytorch.

n = 1

>>> t = torch.tensor([-1.])
>>> torch.polygamma(1, t)
tensor([1.2914e+15])
>>> scipy.special.polygamma(1, t.numpy())
array([inf], dtype=float32)

>>> t = torch.tensor([-501.])
>>> torch.polygamma(1, t)
tensor([2.0831e+09])
>>> scipy.special.polygamma(1, t.numpy())
array([inf], dtype=float32)

>>> t = torch.tensor([-float('inf')])
>>> torch.polygamma(1, t)
tensor([nan])
>>> scipy.special.polygamma(1, t.numpy())
array([inf], dtype=float32)
>>> 

n > 1

>>> t = torch.tensor([float('inf')])
>>> torch.polygamma(2, t)
tensor([nan])
>>> scipy.special.polygamma(2, t.numpy())
array([-0.], dtype=float32)
>>> torch.polygamma(3, t)
tensor([nan])
>>> scipy.special.polygamma(3, t.numpy())
array([0.], dtype=float32)
>>> torch.polygamma(4, t)
tensor([nan])
>>> scipy.special.polygamma(4, t.numpy())
array([-0.], dtype=float32)

cc @mruberry @rgommers @kshitij12345

_Originally posted in #55357

[pearu]

Re

>>> t = torch.tensor([-1.])
>>> torch.polygamma(1, t)
tensor([1.2914e+15])

The problem with polygamma(1, t) is that it internally uses trigamma that does not implement the non-positive-integer arguments check. The reported value is a result of trigamma(t) approximation (pi/sin(pi*t))**2+... that for the given t evaluates to

>>> (torch.pi/torch.sin(t*torch.pi)) ** 2
tensor(1.2914e+15)

It is finite because sin(pi*t) is never exactly zero for a floating point value t.
A fix is to insert std::floor(t) == t check into trigamma that returns inf when the check evaluates true.

Re

>>> t = torch.tensor([-float('inf')])
>>> torch.polygamma(1, t)
tensor([nan])
>>> scipy.special.polygamma(1, t.numpy())
array([inf], dtype=float32)

and other examples with -inf. Here, the torch.polygamma results is as expected while scipy.special.polygamma results are incorrect: polygamma has poles at nonpositive integer values for any n; when evaluating polygamma at -inf, the only reasonable result is nan.