Fix Dirichlet.log_prob() when x=0 and alpha=1

[kalekundert]

Dirichlet.log_prob() incorrectly returns NaN in the case where $x_i=0$ and $\alpha_i=1$. The Dirichlet PDF is given by:
$$\frac{1}{B(\alpha)} \prod_{i=1}^{K} x_i^{\alpha_i - 1}$$
So this corresponds to the case where one of the terms has the form $0^0=1$. The logarithm of such a term should be 0, but you get NaN if you try to calculate it as 0 * log(0).

This PR implements the same algorithm that scipy.stats.dirichlet uses to avoid this behavior, namely xlogy(alpha - 1, x) instead of (alpha - 1) * log(x). It also adds a test case comparing the pytorch and scipy implementations for this specific case.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/103605

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit e001407:
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: kalekundert / name: Kale Kundert (e001407)

[albanD]

Change sounds good!
Just a small question on the test.

[albanD]

Do you really need to override the tolerances here? I would expect the default one to work fine here?

[kalekundert]

I just copied the whole assertion line from the existing Dirichlet.log_prob() test. I don't know if there's a good reason why these tolerances were chosen in the first place, though. Maybe there's a test environment that uses low-precision floats or something? In any case, my goal was to be consistent. Also, this test only really needs to distinguish between NaN and not NaN, so for that the tolerance doesn't matter.

[albanD]

Ok let's keep it as is in the name of consistency then. I agree that it does check what we want here anyways.

[albanD]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here