pymc-devs / pymc-devs/pytensor

Add rewrite for `log(gamma) -> gammaln`

Open
#1,181 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

graph rewriting stabilization
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

https://github.com/pymc-devs/pytensor/blob/911c6a33c2bea6bf1d5b628154e84c43cbed1c63/pytensor/tensor/rewriting/math.py#L3646-L3651

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)

https://github.com/pymc-devs/pytensor/blob/ad55b69f3d13f11c6a9a57823c2a88a966db8b1a/pytensor/tensor/special.py#L799-L804

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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.