Derive logprob for exp2, log2, log10, log1p, expm1, log1mexp, log1pexp (softplus), and sigmoid transformations #6826

LukeLB · 2023-07-13T18:24:47Z

This builds upon the previous pull requests, #6664 and #6775, and completes the work of #6631.

I have attempted to rewrite the logp graph only for log1p, expm1, and log1pexp (softplus). My reasoning is that in these cases, the inputs are transformed directly without affecting the backward or log_jac_det transforms. For all other cases, I have made changes to the existing transform classes.

Although some tests are still failing (see below), I'm uncertain about the reasons for the first two failures. However, it seems that the final failure is due to a floating point error, and I'm unsure how to resolve it.

FAILED test_transforms.py::test_transformed_logprob[at_dist6-dist_params6-<lambda>-size6] - AssertionError: 
E           AssertionError: 
E           Arrays are not almost equal to 4 decimals
E           
E           Mismatched elements: 1 / 1 (100%)
E           Max absolute difference: 0.03384969
E           Max relative difference: 0.00578904
E            x: array(-5.8811)
E            y: array(-5.8472)
FAILED test_transforms.py::test_mixture_transform - assert False
E       assert False
E        +  where False = equal_computations([Sum{axes=None}.0], [Sum{axes=None}.0])
FAILED test_transforms.py::test_check_jac_det[transform9] - AssertionError: 
E           AssertionError: 
E           Not equal to tolerance rtol=1e-07, atol=0
E           
E           Mismatched elements: 1 / 2 (50%)
E           Max absolute difference: 6.59892563e-09
E           Max relative difference: 1.69410928e-07
E            x: array([ 1.05966 , -0.038952])
E            y: array([ 1.05966 , -0.038952])

Currently, there are no specific tests written for the transforms log1p, log1mexp, log1pexp (softplus), and sigmoid. I would appreciate assistance in developing tests for these transforms.
...

Checklist

Explain important implementation details 👆
Make sure that the pre-commit linting/style checks pass.
Link relevant issues (preferably in nice commit messages)
Are the changes covered by tests and docstrings?
Fill out the short summary sections 👇

Major / Breaking Changes

Changes to existing Exp and Log Transform classes

New features

The ability to derive log prob for the following transforms exp2, log2, log10, log1p, expm1, log1mexp, log1pexp (softplus), and sigmoid.

📚 Documentation preview 📚: https://pymc--6826.org.readthedocs.build/en/6826/

codecov · 2023-07-13T18:38:04Z

Codecov Report

Merging #6826 (47f08ea) into main (ddd1d4b) will increase coverage by 0.10%.
Report is 8 commits behind head on main.
The diff coverage is 100.00%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #6826      +/-   ##
==========================================
+ Coverage   92.05%   92.16%   +0.10%     
==========================================
  Files          96      100       +4     
  Lines       16448    16877     +429     
==========================================
+ Hits        15142    15554     +412     
- Misses       1306     1323      +17

Files Changed	Coverage Δ
pymc/logprob/transforms.py	`94.82% <100.00%> (+0.40%)`	⬆️

... and 13 files with indirect coverage changes

ricardoV94

Thanks a lot for picking this up. My suggestions are all about being a bit lazier (less code to maintain and test), but the idea is totally right!

pymc/logprob/transforms.py

ricardoV94 · 2023-07-14T07:57:47Z

My reasoning is that in these cases, the inputs are transformed directly without affecting the backward or log_jac_det transforms. For all other cases, I have made changes to the existing transform classes.

You don't need to worry about jacobian when rewriting into any equivalent forms, not just those. For instance for stuff like log2(c) -> log(c) / log(2), PyMC will have to introduce the right jacobian when deriving the probability of what it sees as a chained scale and log transformations. It doesn't matter what it was before

pymc/logprob/transforms.py

ricardoV94 · 2023-07-16T22:49:07Z

pymc/logprob/transforms.py

+    [inp] = node.inputs
+
+    if isinstance(node.op.scalar_op, Exp2):
+        return [pt.power(2, inp)]


I don't think we support this one, only powers with fixed exponent and variable base

You can convert to exp(ln(2)*x) instead, which PyMC will know how to handle

We should actually use that for any power(const, x) -> exp(log(const) * x) which we currently don't support. But maybe that's better left for another PR?

It requires checking we are interested in x and not const.

Ah yeah I had a feeling that one may be an issue, I'll make the change. I'll add a new function to generalise this power(const, x) -> exp(log(const) * x) functionality.

So I've just pushed this new functionality. Do we care that it won't work for const <= 0?

Let's open a separate issue / PR for this. I thought it over a bit and I think we can definitely support a couple of cases.

power(const, x), for any const > 0 and any x power(const, x), for any const and discrete x (we can play with `log(abs(neg_const))` and x's parity)

The first case we don't have to constrain ourselves to actual "constants", we can add a symbolic assert that const > 0.

The second requires us to implement transforms for discrete variables, which would probably need #6360 first, so we can focus on the first case, which is also probably more useful anyway.

We just have to make sure not to rewrite stuff like power(x, const) accidentally as those are implemented via our PowerTransform. This can be done by checking which of the inputs has a path to unvalued random variables.

pymc/logprob/transforms.py

ricardoV94 · 2023-08-01T15:33:44Z

Rebasing from main should unstuck the tests

review-notebook-app · 2023-08-02T08:53:25Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

ricardoV94 · 2023-08-02T09:24:04Z

I imagine you pull --merge instead of pull --rebase? This makes github show unrelated commits in the PR :/

https://stackoverflow.com/a/71074635

LukeLB · 2023-08-04T15:40:06Z

Woops, sorry about that I will fix

LukeLB · 2023-08-05T18:07:29Z

It seems that the code coverage test is failing because if the new node rewriting functions. Should we be testing them? Or are they fine because they rely on existing functionality thats already tested?

ricardoV94 · 2023-08-07T10:47:17Z

It seems that the code coverage test is failing because if the new node rewriting functions. Should we be testing them? Or are they fine because they rely on existing functionality thats already tested?

We should test them. We can test that the output logp is the same for the new cases we support and what we think they are equivalent to, using the equal_computations helper. Something like

def test():
  base_rv = pt.random.normal(name="base_rv")
  vv = pt.scalar("vv")
  
  logp_test = logp(pt.log1p(base_rv), vv)
  logp_ref = logp(pt.log(1 + base), vv)

  assert equal_computations([logp_test], [logp_ref])

ricardoV94 · 2023-08-24T11:21:03Z

@LukeLB shall we push the PR to the finish line? Let me know if you don't have the time right now

LukeLB · 2023-08-28T20:54:53Z

Hey really sorry about the lack of communication, I've been on holiday. I'm going to have a look at this this week.

ricardoV94 · 2023-08-29T08:02:39Z

No worries, I hope you had a good holiday!

LukeLB · 2023-08-29T14:24:24Z

Ahhh done the same thing as before and somehow merged unrelated commits! Will fix on my side

LukeLB · 2023-08-29T14:27:03Z

While I fix that I am having a problems with two of the tests. For the test

TRANSFORMATIONS = {
    "log1p": (pt.log1p, lambda x: pt.log(1 + x)),
    "softplus": (pt.softplus, lambda x: pt.log(1 + pt.exp(x))),
    "log1mexp": (pt.log1mexp, lambda x: pt.log(1 - pt.exp(pt.neg(x)))),
    "log2": (pt.log2, lambda x: pt.log(x) / pt.log(2)),
    "log10": (pt.log10, lambda x: pt.log(x) / pt.log(10)),
    "exp2": (pt.exp2, lambda x: pt.exp(pt.log(2) * x)),
    "expm1": (pt.expm1, lambda x: pt.exp(x) - 1),
    "sigmoid": (pt.sigmoid, lambda x: 1 / (1 + pt.exp(-x))),
}


@pytest.mark.parametrize("transform", TRANSFORMATIONS.keys())
def test_special_log_exp_transforms(transform):
    base_rv = pt.random.normal(name="base_rv")
    vv = pt.scalar("vv")

    transform_func, ref_func = TRANSFORMATIONS[transform]
    transformed_rv = transform_func(base_rv)
    ref_transformed_rv = ref_func(base_rv)

    logp_test = logp(transformed_rv, vv)
    logp_ref = logp(ref_transformed_rv, vv)

    assert equal_computations([logp_test], [logp_ref])

when transform is log2 or log10 then test fails for equal computation, I'm not sure what is causing that...

ricardoV94 · 2023-09-04T15:41:37Z

While I fix that I am having a problems with two of the tests. For the test

TRANSFORMATIONS = {
    "log1p": (pt.log1p, lambda x: pt.log(1 + x)),
    "softplus": (pt.softplus, lambda x: pt.log(1 + pt.exp(x))),
    "log1mexp": (pt.log1mexp, lambda x: pt.log(1 - pt.exp(pt.neg(x)))),
    "log2": (pt.log2, lambda x: pt.log(x) / pt.log(2)),
    "log10": (pt.log10, lambda x: pt.log(x) / pt.log(10)),
    "exp2": (pt.exp2, lambda x: pt.exp(pt.log(2) * x)),
    "expm1": (pt.expm1, lambda x: pt.exp(x) - 1),
    "sigmoid": (pt.sigmoid, lambda x: 1 / (1 + pt.exp(-x))),
}


@pytest.mark.parametrize("transform", TRANSFORMATIONS.keys())
def test_special_log_exp_transforms(transform):
    base_rv = pt.random.normal(name="base_rv")
    vv = pt.scalar("vv")

    transform_func, ref_func = TRANSFORMATIONS[transform]
    transformed_rv = transform_func(base_rv)
    ref_transformed_rv = ref_func(base_rv)

    logp_test = logp(transformed_rv, vv)
    logp_ref = logp(ref_transformed_rv, vv)

    assert equal_computations([logp_test], [logp_ref])

when transform is log2 or log10 then test fails for equal computation, I'm not sure what is causing that...

Try to look at pytensor.dprint(logp_test) and pytensor.dprint(logp_ref) and see if anything obvious comes across. Even something like x + 2 vs 2 + x would make that function return False

LukeLB · 2023-09-04T16:58:13Z

@ricardoV94 OK I think I've got to the bottom of what is causing the failure. Looks like a floating point in precision in one of the nodes of the graph between the test case and the reference case.

However I'm not sure I can control that in the original lambda function that I've written. Any ideas?

ricardoV94 · 2023-09-04T17:04:21Z

However I'm not sure I can control that in the original lambda function that I've written. Any ideas?

Hmm. Let's do a logp evaluation for that one (separate test) and check for output closeness?

Co-authored-by: Ricardo Vieira <[email protected]>

…ion in the graph)

ricardoV94 · 2023-09-05T13:19:49Z

File changes is showing some accidental overwriting of previous changes? https://github.com/pymc-devs/pymc/pull/6826/files

tests/logprob/test_transforms.py

pymc/logprob/transforms.py

tests/logprob/test_transforms.py

…sion

ricardoV94 · 2023-09-07T14:01:21Z

Awesome work @LukeLB. Do you want to pursue the power exponent next?

LukeLB · 2023-09-07T21:02:24Z

Thanks @ricardoV94, I feel like I learnt a lot on this one!

Do you want to pursue the power exponent next?

Yep I'll start looking into that, if I run into any trouble I'll drop you a mesage on slack.

twiecki · 2023-09-14T07:41:33Z

Congrats @LukeLB! This is a major and non-trivial contribution.

ricardoV94 reviewed Jul 14, 2023

View reviewed changes

pymc/logprob/transforms.py Outdated Show resolved Hide resolved

pymc/logprob/transforms.py Outdated Show resolved Hide resolved

pymc/logprob/transforms.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Jul 16, 2023

View reviewed changes

ricardoV94 reviewed Jul 17, 2023

View reviewed changes

pymc/logprob/transforms.py Outdated Show resolved Hide resolved

LukeLB force-pushed the issue_#6631_chapter3_the_end branch from 0f50a55 to 0cfd1f2 Compare August 5, 2023 11:05

ricardoV94 added enhancements logprob labels Aug 7, 2023

ricardoV94 changed the title ~~[WIP] Derive logprob for exp2, log2, log10, log1p, expm1, log1mexp, log1pexp (softplus), and sigmoid transformations~~ Derive logprob for exp2, log2, log10, log1p, expm1, log1mexp, log1pexp (softplus), and sigmoid transformations Aug 7, 2023

Luke LB added 7 commits September 4, 2023 21:51

LogTransform with arb base n now works

6f651e6

log2 and log10 now working

e9ce641

log1p, softplus, and log1mexp now working

5035ecd

merge conflicts in tests_transforms

5f4001a

3 tests failing

85e0f37

implemented rewrites and reverted Transform classes

f4bf24b

removed additions to node rewriter

09f15fb

Luke LB and others added 14 commits September 4, 2023 21:51

reverted back to just allowing exp2

588e89b

removed measurable rewrite

1168133

written test place holder

1748588

tests running but log2, log10, and expm1 failing

1b3b6ed

expm1 test now passing

59a37fc

LogTransform with arb base n now works

f7c5938

log2 and log10 now working

c47ec86

log1p, softplus, and log1mexp now working

1347695

updated tests but failing

117eff4

3 tests failing

fc5eed7

implemented rewrites and reverted Transform classes

27e11a5

removed additions to node rewriter

18901ea

Update pymc/logprob/transforms.py

9412b39

Co-authored-by: Ricardo Vieira <[email protected]>

fixed test failing for log2 and log10 (due to floating point inprecis…

a912b20

…ion in the graph)

LukeLB force-pushed the issue_#6631_chapter3_the_end branch from 17123ac to a912b20 Compare September 4, 2023 21:05

accidently committed notebook

86327f3

fixed removed code from rebase

346f085

ricardoV94 reviewed Sep 6, 2023

View reviewed changes

tests/logprob/test_transforms.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Sep 6, 2023

View reviewed changes

pymc/logprob/transforms.py Outdated Show resolved Hide resolved

moved test back into TestClass after rebase

2882a0e

ricardoV94 reviewed Sep 7, 2023

View reviewed changes

tests/logprob/test_transforms.py Outdated Show resolved Hide resolved

fixed duplicat test and removed negation from log1mexp logprob conver…

47f08ea

…sion

ricardoV94 approved these changes Sep 7, 2023

View reviewed changes

ricardoV94 merged commit 39a2975 into pymc-devs:main Sep 7, 2023
21 checks passed

This was referenced Sep 7, 2023

Support logp derivation of power exponent #6896

Closed

Implement more types of measurable elemwise transformations #6631

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Derive logprob for exp2, log2, log10, log1p, expm1, log1mexp, log1pexp (softplus), and sigmoid transformations #6826

Derive logprob for exp2, log2, log10, log1p, expm1, log1mexp, log1pexp (softplus), and sigmoid transformations #6826

LukeLB commented Jul 13, 2023 •

edited by github-actions bot

Loading

codecov bot commented Jul 13, 2023 •

edited

Loading

ricardoV94 left a comment

ricardoV94 commented Jul 14, 2023 •

edited

Loading

ricardoV94 Jul 16, 2023 •

edited

Loading

ricardoV94 Jul 16, 2023 •

edited

Loading

ricardoV94 Jul 16, 2023 •

edited

Loading

LukeLB Jul 17, 2023

LukeLB Jul 17, 2023 •

edited

Loading

ricardoV94 Aug 7, 2023 •

edited

Loading

ricardoV94 commented Aug 1, 2023

review-notebook-app bot commented Aug 2, 2023

ricardoV94 commented Aug 2, 2023

LukeLB commented Aug 4, 2023

LukeLB commented Aug 5, 2023

ricardoV94 commented Aug 7, 2023 •

edited

Loading

ricardoV94 commented Aug 24, 2023

LukeLB commented Aug 28, 2023

ricardoV94 commented Aug 29, 2023

LukeLB commented Aug 29, 2023

LukeLB commented Aug 29, 2023

ricardoV94 commented Sep 4, 2023 •

edited

Loading

LukeLB commented Sep 4, 2023

ricardoV94 commented Sep 4, 2023

ricardoV94 commented Sep 5, 2023 •

edited

Loading

ricardoV94 commented Sep 7, 2023

LukeLB commented Sep 7, 2023

twiecki commented Sep 14, 2023

Derive logprob for exp2, log2, log10, log1p, expm1, log1mexp, log1pexp (softplus), and sigmoid transformations #6826

Derive logprob for exp2, log2, log10, log1p, expm1, log1mexp, log1pexp (softplus), and sigmoid transformations #6826

Conversation

LukeLB commented Jul 13, 2023 • edited by github-actions bot Loading

Major / Breaking Changes

New features

codecov bot commented Jul 13, 2023 • edited Loading

Codecov Report

ricardoV94 left a comment

Choose a reason for hiding this comment

ricardoV94 commented Jul 14, 2023 • edited Loading

ricardoV94 Jul 16, 2023 • edited Loading

Choose a reason for hiding this comment

ricardoV94 Jul 16, 2023 • edited Loading

Choose a reason for hiding this comment

ricardoV94 Jul 16, 2023 • edited Loading

Choose a reason for hiding this comment

LukeLB Jul 17, 2023

Choose a reason for hiding this comment

LukeLB Jul 17, 2023 • edited Loading

Choose a reason for hiding this comment

ricardoV94 Aug 7, 2023 • edited Loading

Choose a reason for hiding this comment

ricardoV94 commented Aug 1, 2023

review-notebook-app bot commented Aug 2, 2023

ricardoV94 commented Aug 2, 2023

LukeLB commented Aug 4, 2023

LukeLB commented Aug 5, 2023

ricardoV94 commented Aug 7, 2023 • edited Loading

ricardoV94 commented Aug 24, 2023

LukeLB commented Aug 28, 2023

ricardoV94 commented Aug 29, 2023

LukeLB commented Aug 29, 2023

LukeLB commented Aug 29, 2023

ricardoV94 commented Sep 4, 2023 • edited Loading

LukeLB commented Sep 4, 2023

ricardoV94 commented Sep 4, 2023

ricardoV94 commented Sep 5, 2023 • edited Loading

ricardoV94 commented Sep 7, 2023

LukeLB commented Sep 7, 2023

twiecki commented Sep 14, 2023

LukeLB commented Jul 13, 2023 •

edited by github-actions bot

Loading

codecov bot commented Jul 13, 2023 •

edited

Loading

ricardoV94 commented Jul 14, 2023 •

edited

Loading

ricardoV94 Jul 16, 2023 •

edited

Loading

ricardoV94 Jul 16, 2023 •

edited

Loading

ricardoV94 Jul 16, 2023 •

edited

Loading

LukeLB Jul 17, 2023 •

edited

Loading

ricardoV94 Aug 7, 2023 •

edited

Loading

ricardoV94 commented Aug 7, 2023 •

edited

Loading

ricardoV94 commented Sep 4, 2023 •

edited

Loading

ricardoV94 commented Sep 5, 2023 •

edited

Loading