YouTipPytorch Torch Polygamma
## PyTorch torch.polygamma
The `torch.polygamma` function in PyTorch is used to compute the element-wise $n^{th}$ derivative of the digamma function (which is the logarithmic derivative of the Gamma function) for a given input tensor.
---
## Mathematical Definition
The polygamma function of order $n$ (denoted as $\psi^{(n)}(x)$) is mathematically defined as:
$$\psi^{(n)}(x) = \frac{d^{n+1}}{dx^{n+1}} \ln \Gamma(x)$$
Where:
* $\Gamma(x)$ is the Gamma function.
* $\psi^{(0)}(x)$ is the digamma function.
* $\psi^{(1)}(x)$ is the trigamma function, and so on.
---
## Syntax
```python
torch.polygamma(n, input, *, out=None) -> Tensor
```
### Parameters
| Parameter | Type | Description |
| :--- | :--- | :--- |
| **`n`** | `int` | The order of the polygamma function to compute. Must be a non-negative integer ($n \ge 0$). |
| **`input`** | `Tensor` | The input tensor containing real values. |
| **`out`** | `Tensor` *(optional)* | The output tensor where the result will be written. |
### Return Value
* Returns a new tensor with the same shape as `input` containing the computed polygamma values.
---
## Code Examples
### Example 1: Basic Usage (Trigamma Function, $n = 1$)
The following example demonstrates how to compute the first polygamma function (trigamma function) for a 1D tensor.
```python
import torch
# Create an input tensor
x = torch.tensor([1.0, 2.0, 3.0])
# Define the order of the polygamma function (n = 1 represents trigamma)
n = 1
# Compute the polygamma function
result = torch.polygamma(n, x)
print("Input tensor:")
print(x)
print("\nResult (n=1):")
print(result)
```
**Output:**
```text
Input tensor:
tensor([1., 2., 3.])
Result (n=1):
tensor([1.6449, 0.6449, 0.3949])
```
---
### Example 2: Higher-Order Polygamma ($n = 2$)
You can compute higher-order derivatives by increasing the value of `n`. Here, we compute the second polygamma function (tetragamma function).
```python
import torch
# Create a 2D input tensor
x = torch.tensor([[0.5, 1.5], [2.5, 3.5]])
# Compute the second polygamma function (n = 2)
result = torch.polygamma(2, x)
print("Result (n=2):")
print(result)
```
**Output:**
```text
Result (n=2):
tensor([[-16.8288, -1.9022],
[ -0.4022, -0.1511]])
```
---
### Example 3: In-place Operation
PyTorch also provides an in-place version of the function: `torch.polygamma_()`. This modifies the input tensor directly without allocating new memory.
```python
import torch
x = torch.tensor([1.0, 2.0, 3.0])
# Perform in-place polygamma computation (n = 1)
x.polygamma_(1)
print("In-place result:")
print(x)
```
**Output:**
```text
In-place result:
tensor([1.6449, 0.6449, 0.3949])
```
---
## Technical Considerations
1. **Domain Restrictions**: The polygamma function is defined for all real numbers except for non-positive integers ($0, -1, -2, \dots$), where the function poles exist and the values approach infinity. Passing these values will result in `NaN` or `inf`.
2. **Data Types**: The input tensor must contain floating-point or complex data types. If integer tensors are passed, PyTorch will automatically promote them to floating-point types.
3. **Gradient Computation**: `torch.polygamma` fully supports autograd and backpropagation, making it suitable for use in custom loss functions or neural network layers.