pymc-devs / pymc-devs/pytensor
Add rewrite for `log(gamma) -> gammaln`
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 644
- Forks
- 208
- Avg merge
- 2d 14h
- Merged PRs (30d)
- 16
Description
Description
We're missing this simple rewrite:
import pytensor.tensor as pt
from pytensor.graph import rewrite_graph
x = pt.scalar("x")
out = pt.log(pt.gamma(x))
new_out = rewrite_graph(out, include=("canonicalize", "stabilize", "specialize"))
new_out.dprint()
Can be done easily with PatternNodeRewriter as in
We could also add rewrites for common combinatorics expressions like
naive_betaln = pt.log((pt.gamma(x) * pt.gamma(y)) / pt.gamma(x + y)
betaln = pt.gammaln(x) + pt.gammaln(y) - pt.gammaln(x + y)
Or for log(poch): https://github.com/pymc-devs/pytensor/blob/ad55b69f3d13f11c6a9a57823c2a88a966db8b1a/pytensor/tensor/special.py#L767-L772
For these more general cases we can probably use something more flexible than the PatternNodeRewriter. We want to apply as long as we know all the terms inside are factorials/gammas/exps (positive things that easily blow up). This is a narrow/easier subset of https://github.com/pymc-devs/pytensor/discussions/177
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start in pytensor/tensor/rewriting/math.py at the PatternNodeRewriter example around lines 3646-3651, then inspect the related gammaln, betaln, and poch definitions in pytensor/tensor/special.py. Run the issue's rewrite_graph example and implement the log(gamma) to gammaln rewrite; done means the rewritten graph uses gammaln for this expression.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers
- Issue type
- Feature
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100