Check if logprob derivation works with specialized `log1p`, `log1mexp`, `expm1`

[ricardoV94]

Description

If pm.logp(pm.log1p(pm.HalfNormal.dist()) does not work, we should canonicalize into the form that the logprob submodule can understand as done in #6371

[tamastokes]

Hi,

Just to confirm that I understand this problem right: I did try and got a NotImplementedError for pm.logp(pt.log1p(pm.HalfNormal.dist()), 0.5), see below. (Note that I used pt.log1p, as I haven't found pm.log1p.). So, as I understand, this task would be to:

1. Find out what's missing for pm.logp() to be able to deal with Elemwise{log1p,no_inplace}.
2. Then,

• either convert the Elemwise{log1p,no_inplace} node to a form that pm.logp() deals with (first guess would be to undo the log1p rewrite to sub(log,1)), however, in this case my question would be how we can prevent the rewrite to take place again?)
• or add logic to pm.logp() for it to be able to deal with the log1p.

4. Same for log1mexp and expm1.

Is my understanding correct?

Here's the Exception I got:

>>> import pymc as pm
>>> from pytensor import tensor as pt
>>> lp = pm.logp(pt.log1p(pm.HalfNormal.dist()), 0.5)
Traceback (most recent call last):
  File "/Users/tokes/src/pymc/pymc/distributions/logprob.py", line 177, in logp
    return logp_logprob(rv, value)
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/tokes/src/pymc/pymc/logprob/abstract.py", line 53, in logprob
    logprob = _logprob(rv_var.owner.op, rv_values, *rv_var.owner.inputs, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Caskroom/miniforge/base/envs/pymc-dev/lib/python3.11/functools.py", line 909, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/tokes/src/pymc/pymc/logprob/abstract.py", line 94, in _logprob
    raise NotImplementedError(f"Logprob method not implemented for {op}")
NotImplementedError: Logprob method not implemented for Elemwise{log1p,no_inplace}

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/Users/tokes/src/pymc/pymc/distributions/logprob.py", line 186, in logp
    return factorized_joint_logprob({rv: value}, warn_missing_rvs=False)[value]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/tokes/src/pymc/pymc/logprob/joint_logprob.py", line 220, in factorized_joint_logprob
    raise RuntimeError(
RuntimeError: The logprob terms of the following value variables could not be derived: {TensorConstant{0.5}}

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/tokes/src/pymc/pymc/distributions/logprob.py", line 188, in logp
    raise NotImplementedError("PyMC could not infer logp of input variable.") from exc
NotImplementedError: PyMC could not infer logp of input variable.
>>> 

[ricardoV94]

Yes it makes sense to rewrite it to the form PyMC already understands. Note that log1p(x) is log(add(1, x))

We should restrict the rewrite to the cases where x is a measurable variable without value variable before. This is done by some of the intermediate rewrites in logprob.transform.py already.

https://github.com/pymc-devs/pymc/blob/main/pymc/logprob/transforms.py#L406-L426

[ricardoV94]

Closing in favor of #6631