Differentiating beta distribution PDF w.r.t. shape parameters

[MatthewGerber]

Hello -

I am trying to use Jax to calculate the derivative of the beta distribution's PDF w.r.t. its shape parameters. The output does not match my expectations. Here's my minimal example, with my questions and output embedded and the plot attached. Much appreciation in advance to anyone who can shed light on these questions.

Matt

import jax.scipy.stats as jstats
from jax import grad
import matplotlib.pyplot as plt
from scipy import stats

# beta distribution parameters
a = 0.1
b = 0.1
loc = 0.0
scale = 1.0

# sample the beta distribution and plot histogram
fig, ax = plt.subplots(1, 1)
r = stats.beta.rvs(a=a, b=b, size=1000, loc=loc, scale=scale)
ax.hist(r, density=True, histtype='stepfilled', alpha=0.2)
ax.set_xlim(-0.5, 1.5)
plt.show()

# what happens to the beta pdf at x=0.01 when we increase a? that is, what is the gradient of the pdf at x=0.01 wrt a?
x = 0.01
print(f'Beta PDF at x:  {stats.beta.pdf(x=x, a=a, b=b, loc=loc, scale=scale)}')  # 3.229531203997695
print(f'Beta PDF at x, with larger a:  {stats.beta.pdf(x=x, a=a+1, b=b, loc=loc, scale=scale)}')  # 0.0645906240799539

# answer:  the beta pdf at x=0.01 decreases if we increase a. that is, the gradient of the pdf at x=0.01 wrt a is negative.

# can we answer the same question with jax?

def jax_beta_pdf(
        a: float
):
    return jstats.beta.pdf(x=x, a=a, b=b, loc=loc, scale=scale)

jax_beta_pdf_grad = grad(jax_beta_pdf, argnums=0)

a_grad = jax_beta_pdf_grad(a)

print(f'{a_grad}')  # 1.710184097290039 -- what does this mean? shouldn't this be negative?

[sharadmv]

This is tricky!

# answer:  the beta pdf at x=0.01 decreases if we increase a.

When a = 0.1 and b=0.1, the beta has modes at 0 and 1 (you can see the pdf shoot towards infinity in those regions). By increasing a slightly, we actually increase the density of x = 0.01 until a gets too big and then the density of x = 0.01 starts decreasing again. We can see this in the plot of jax_beta_pdf for various values of a:

I hope that answered your question!

[MatthewGerber]

Ah! Yes, that was tricky. I had unfortunately been playing around with a = 0.1, which is quite close to the peak of your graph as shown in my version below:

So although increasing a by 1 results in a net decrease in the PDF (second vertical line), the gradient at x = 0.01 w.r.t. a was still positive. Many thanks.