Implement `default_transform` and `transform` argument for distributions #7207

aerubanov · 2024-03-21T18:31:38Z

Description

Related Issue

Closes Implement default_transform and transform argument for distributions #5674
Related to #

Type of change

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

ricardoV94

Thanks a lot, left some small user-quality of life suggestions

pymc/model/core.py

codecov · 2024-03-22T13:33:55Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 92.28%. Comparing base (034b9a4) to head (4c94e9e).
Report is 19 commits behind head on main.

❗ Current head 4c94e9e differs from pull request most recent head 704aac6. Consider uploading reports for the commit 704aac6 to get more accurate results

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #7207      +/-   ##
==========================================
- Coverage   92.29%   92.28%   -0.01%     
==========================================
  Files         101      100       -1     
  Lines       16892    16906      +14     
==========================================
+ Hits        15590    15602      +12     
- Misses       1302     1304       +2

Files	Coverage Δ
pymc/distributions/distribution.py	`93.96% <100.00%> (+0.03%)`	⬆️
pymc/model/core.py	`92.26% <100.00%> (+0.44%)`	⬆️
pymc/model/fgraph.py	`97.39% <100.00%> (-0.51%)`	⬇️

... and 16 files with indirect coverage changes

pymc/distributions/distribution.py

aerubanov · 2024-03-27T08:03:36Z

Looks like I fixed all failed test cases. Going add some new tests and changes to documentation as a next steps.

ricardoV94 · 2024-04-02T09:25:38Z

Looks like I fixed all failed test cases. Going add some new tests and changes to documentation as a next steps.

@mkusnetsov took the documentation initiative in #7232 so we should be good to go just with docstrings and tests

Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com>

aerubanov · 2024-04-09T16:56:34Z

@ricardoV94 I added some test cases to check warning it transform=None and transform order if both default_transform and transform are using

ricardoV94 · 2024-04-15T08:44:44Z

pymc/distributions/distribution.py

@@ -397,6 +398,15 @@ def __new__(
        if not isinstance(name, string_types):
            raise TypeError(f"Name needs to be a string but got: {name}")

+        if transform is None and default_transform is UNSET:


This warning should be in the relevant section in pm.Model instead of Distribution

pymc/model/core.py

ricardoV94 · 2024-04-15T08:50:28Z

tests/distributions/test_mixture.py

-                Mixture("mix4", w=[0.5, 0.5], comp_dists=comp_dists, transform=None)
+                Mixture("mix4", w=[0.5, 0.5], comp_dists=comp_dists, default_transform=None)
+
+            with pytest.warns(


This functionality shouldn't be tested here, since it's not specific to Mixture. Probably in model/test_core.py there should be related stuff already? I see you already have it there, so is this test needed?

tests/distributions/test_transform.py

ricardoV94 · 2024-04-15T08:52:38Z

tests/logprob/test_utils.py

+            x = pm.Uniform("x", lower=0, upper=1, transform=transform, default_transform=None)
            # Operation between the variables provides a regression test for #7054
-            y = pm.Uniform("y", lower=0, upper=pt.exp(x), transform=transform)
-            z = pm.Uniform("z", lower=0, upper=y, transform=transform)
-            w = pm.Uniform("w", lower=0, upper=pt.square(z), transform=transform)
+            y = pm.Uniform(
+                "y", lower=0, upper=pt.exp(x), transform=transform, default_transform=None
+            )
+            z = pm.Uniform("z", lower=0, upper=y, transform=transform, default_transform=None)
+            w = pm.Uniform(
+                "w", lower=0, upper=pt.square(z), transform=transform, default_transform=None


Pass transform to default_transform

ricardoV94 · 2024-04-15T08:55:57Z

tests/model/test_core.py

+
+        with pm.Model() as model:
+            x = pm.Normal("x", transform=DummyTransform(2), default_transform=DummyTransform(1))
+        assert transform_order == [1, 2]


You can use regular transforms, and simply assert the obtained transform is a Chain, which as a property like transform_list that includes the transforms, and you can assert those are also the expected ones. The transform is available in models.rvs_to_transforms[x]

Also, would be nice to include a numerical example that would have led to nan or -inf probability before the change, like an ordered mixture of LogNormals evaluated at -1

Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com>

ricardoV94

Besides making the test_transform_order simpler, just need some tweaks to the docstrings. They don't describe the parameter. We can also give better type hint than Any. If mypy complaints just revert to Any

ricardoV94 · 2024-04-16T08:32:52Z

pymc/model/core.py

+        total_size=None,
+        dims=None,
+        transform=UNSET,
+        default_transform=UNSET,


Missing description in the docstrings

ricardoV94 · 2024-04-16T08:33:08Z

pymc/model/core.py

@@ -1288,6 +1299,7 @@ def make_obs_var(
        data: np.ndarray,
        dims,
        transform: Any | None,
+        default_transform: Any | None,


Missing description in the docstrings.

Also type hint isn't great, should be RVTransform | None (or something like that, don't remember the exact class now)

ricardoV94 · 2024-04-16T08:34:26Z

pymc/model/core.py

+        rv_var: TensorVariable,
+        *,
+        transform: Any,
+        default_transform: Any,


Update type hints and docstring

aerubanov · 2024-04-17T16:28:35Z

@ricardoV94 Is there any way to get logp value after transform applied? I`m trying something like

    with pm.Model() as model:
        x1 = pm.LogNormal("x1", 0.0, 1.0)
        x2 = pm.LogNormal("x2", 0.0, 1.0, default_transform=LogTransform())
        assert pm.logp(x1, 1.0).eval() != pm.logp(x2, 1.0).eval()

and got assertion error

ricardoV94 · 2024-04-17T17:20:49Z

@ricardoV94 Is there any way to get logp value after transform applied? I`m trying something like

    with pm.Model() as model:
        x1 = pm.LogNormal("x1", 0.0, 1.0)
        x2 = pm.LogNormal("x2", 0.0, 1.0, default_transform=LogTransform())
        assert pm.logp(x1, 1.0).eval() != pm.logp(x2, 1.0).eval()

and got assertion error

You want to use model.compile_logp(), and you'll have to create two models, one with transform and one without

aerubanov · 2024-04-18T09:51:17Z

@ricardoV94 I added new test case with numerical example for transform args. Please let me know what do you thunk about it.

ricardoV94

Thanks @aerubanov

I left some minor comments, I think this is the last round!

tests/model/test_core.py

ricardoV94 · 2024-04-18T10:26:32Z

tests/model/test_core.py

+        with pm.Model() as model1:
+            x1 = pm.LogNormal("x1", 0, 1, transform=Interval(-2, 2), default_transform=None)
+        with pm.Model() as model3:
+            x2 = pm.LogNormal("x2", 0, 1, transform=Interval(-2, 2), default_transform=log)


This is fine but is not very realistic. What about ordered transform + log which is the example that motivated this PR?

Suggested change

x2 = pm.LogNormal("x2", 0, 1, transform=Interval(-2, 2), default_transform=log)

x2 = pm.LogNormal("x2", 0, 1, transform=Interval(-2, 2), default_transform=log)

The rest of the logic of the test is spot-on!

pymc/model/core.py

ricardoV94 · 2024-04-18T10:27:27Z

pymc/model/core.py

@@ -1230,7 +1239,9 @@ def register_rv(
        dims : tuple
            Dimension names for the variable.
        transform
-            A transform for the random variable in log-likelihood space.
+            Additianal transform which may be applied after default transform.
+        default_transform


Nitpick: show default_transform before transform (also in the signature)?

ricardoV94 · 2024-04-18T10:27:38Z

pymc/model/core.py

+        transform
+            Additianal transform which may be applied after default transform.
+        default_transform


Same as above

ricardoV94 · 2024-04-18T10:27:49Z

pymc/model/core.py

+        transform: Transform
+            Additianal transform which may be applied after default transform.
+
+        default_transform: Transform
+            A transform for the random variable in log-likelihood space.


Same as above

Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com>

aerubanov · 2024-04-18T15:28:43Z

@ricardoV94 please review again

tests/model/test_core.py

ricardoV94 · 2024-04-18T16:08:53Z

tests/model/test_core.py

+    def test_transform_order(self):
+        with pm.Model() as model:
+            x = pm.Normal("x", transform=Interval(0, 1), default_transform=log)
+        assert isinstance(model.rvs_to_transforms[x], ChainedTransform)


Nitpick (feel free to ignore): Save the transform in a separate variable so you don't need to write 3 times model.rvs_to_transforms[x]

Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com>

ricardoV94

Thanks, looks great!

ricardoV94 · 2024-04-18T16:49:26Z

We should follow up with an issue to provide informative warnings when doing prior/posterior predictive sampling of variables with custom non-default transforms, like we do with Potentials

aerubanov marked this pull request as draft March 21, 2024 18:31

ricardoV94 reviewed Mar 22, 2024

View reviewed changes

pymc/model/core.py Show resolved Hide resolved

pymc/model/core.py Outdated Show resolved Hide resolved

pymc/model/core.py Show resolved Hide resolved

pymc/model/core.py Show resolved Hide resolved

ricardoV94 added the enhancements label Mar 22, 2024

ricardoV94 reviewed Mar 25, 2024

View reviewed changes

pymc/distributions/distribution.py Outdated Show resolved Hide resolved

mkusnetsov mentioned this pull request Mar 29, 2024

Add introduction section to Transformations API page to document role of Transforms #7232

Merged

11 tasks

aerubanov and others added 9 commits April 9, 2024 19:08

add default transform argument

8a50f9f

Apply suggestions from code review

7434b9a

Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com>

fix arguments passing

4fde32f

add support for transform=None

5c7519f

fix failed tests

1dc0fcb

Update pymc/distributions/distribution.py

55b0edb

Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com>

fix failed tests

d791ac7

fix test

def0f6e

add tests for transfor args

7d6ecf9

aerubanov force-pushed the 5674-default-transform-arg branch from 1225764 to 7d6ecf9 Compare April 9, 2024 16:11

fix formating

704aac6

aerubanov marked this pull request as ready for review April 9, 2024 16:54

aerubanov requested a review from ricardoV94 April 9, 2024 16:54

ricardoV94 requested changes Apr 15, 2024

View reviewed changes

ricardoV94 changed the title ~~Implement default_transform and transform argument for distributions~~ Implement default_transform and transform argument for distributions Apr 15, 2024

ricardoV94 added the request discussion label Apr 15, 2024

aerubanov and others added 5 commits April 15, 2024 20:17

Update pymc/model/core.py

23bd69f

Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com>

Update tests/distributions/test_transform.py

35b44fe

Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com>

Update tests/distributions/test_transform.py

7f0d1d3

Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com>

pass transform to default transform

80641f8

remove duplicated test case

80fa510

move transform=None warning to core module

655a364

ricardoV94 reviewed Apr 16, 2024

View reviewed changes

aerubanov added 2 commits April 17, 2024 18:50

update type hints and docstrings

995ec93

change test for transforms order

3eefa8e

add new test case

8a005fe

ricardoV94 reviewed Apr 18, 2024

View reviewed changes

aerubanov and others added 3 commits April 18, 2024 16:13

Update pymc/model/core.py

8a708f7

Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com>

change test casses for transform args

412ef7f

fix docstring and change argument order

48f5f85

aerubanov requested a review from ricardoV94 April 18, 2024 15:28

ricardoV94 reviewed Apr 18, 2024

View reviewed changes

tests/model/test_core.py Show resolved Hide resolved

tests/model/test_core.py Outdated Show resolved Hide resolved

tests/model/test_core.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Apr 18, 2024

View reviewed changes

aerubanov and others added 2 commits April 18, 2024 20:24

Apply suggestions from code review

bff4371

Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com>

refactor test case

e0f26b0

aerubanov requested a review from ricardoV94 April 18, 2024 16:27

ricardoV94 approved these changes Apr 18, 2024

View reviewed changes

ricardoV94 merged commit 23e418d into pymc-devs:main Apr 19, 2024
21 checks passed

ricardoV94 added the major Include in major changes release notes section label Apr 19, 2024

aerubanov deleted the 5674-default-transform-arg branch April 19, 2024 07:35

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Implement `default_transform` and `transform` argument for distributions #7207

Implement `default_transform` and `transform` argument for distributions #7207

aerubanov commented Mar 21, 2024 •

edited by github-actions bot

Loading

ricardoV94 left a comment

codecov bot commented Mar 22, 2024 •

edited

Loading

aerubanov commented Mar 27, 2024

ricardoV94 commented Apr 2, 2024

aerubanov commented Apr 9, 2024

ricardoV94 Apr 15, 2024 •

edited

Loading

ricardoV94 Apr 15, 2024 •

edited

Loading

ricardoV94 Apr 15, 2024

ricardoV94 Apr 15, 2024

ricardoV94 left a comment

ricardoV94 Apr 16, 2024

ricardoV94 Apr 16, 2024

ricardoV94 Apr 16, 2024

ricardoV94 Apr 16, 2024

aerubanov commented Apr 17, 2024

ricardoV94 commented Apr 17, 2024

aerubanov commented Apr 18, 2024

ricardoV94 left a comment

ricardoV94 Apr 18, 2024

ricardoV94 Apr 18, 2024

ricardoV94 Apr 18, 2024

ricardoV94 Apr 18, 2024

ricardoV94 Apr 18, 2024

aerubanov commented Apr 18, 2024

ricardoV94 Apr 18, 2024

ricardoV94 left a comment

ricardoV94 commented Apr 18, 2024 •

edited

Loading

	x2 = pm.LogNormal("x2", 0, 1, transform=Interval(-2, 2), default_transform=log)
	x2 = pm.LogNormal("x2", 0, 1, transform=Interval(-2, 2), default_transform=log)

Implement default_transform and transform argument for distributions #7207

Implement default_transform and transform argument for distributions #7207

Conversation

aerubanov commented Mar 21, 2024 • edited by github-actions bot Loading

Description

Related Issue

Type of change

ricardoV94 left a comment

Choose a reason for hiding this comment

codecov bot commented Mar 22, 2024 • edited Loading

Codecov Report

aerubanov commented Mar 27, 2024

ricardoV94 commented Apr 2, 2024

aerubanov commented Apr 9, 2024

ricardoV94 Apr 15, 2024 • edited Loading

Choose a reason for hiding this comment

ricardoV94 Apr 15, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ricardoV94 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

aerubanov commented Apr 17, 2024

ricardoV94 commented Apr 17, 2024

aerubanov commented Apr 18, 2024

ricardoV94 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

aerubanov commented Apr 18, 2024

Choose a reason for hiding this comment

ricardoV94 left a comment

Choose a reason for hiding this comment

ricardoV94 commented Apr 18, 2024 • edited Loading

Implement `default_transform` and `transform` argument for distributions #7207

Implement `default_transform` and `transform` argument for distributions #7207

aerubanov commented Mar 21, 2024 •

edited by github-actions bot

Loading

codecov bot commented Mar 22, 2024 •

edited

Loading

ricardoV94 Apr 15, 2024 •

edited

Loading

ricardoV94 Apr 15, 2024 •

edited

Loading

ricardoV94 commented Apr 18, 2024 •

edited

Loading