Simplify Model new and metaclass #7473

thomasaarholt · 2024-08-23T09:20:31Z

Description

This PR is a refactor that improves the code that creates models and keeps track of the model contexts that are active at a given moment.

I have previously done some work (#6809) on type hinting the Model class. It was a bit difficult to understand what was going on, and the code was lacking documentation that would help further development. In particular, the code creating the Model class and instance, and keeping track of active contexts was rather hacky.

I am on parental leave, which (maybe ironically) gives me the time to take a look behind the scenes and hopefully make some improvements that give back to the community. I do recognize that this PR comes a bit "unannounced", and since it doesn't fix any major issues or introduce hot features, I don't have any expectations for it to be reviewed soon (or even at all).

Each commit is a logical change with a reasonably descriptive text.

This passes all tests in tests/model.

Checklist

Checked that the pre-commit linting/style checks pass
Included tests that prove the fix is effective or that the new feature works
Added necessary documentation (docstrings and/or example notebooks)
If you are a pro: each commit corresponds to a relevant logical change

Type of change

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

thomasaarholt · 2024-08-23T12:52:51Z

The mypy failure is because my code change allows mypy to infer the type of model in pymc/data.py:446 as Model rather than Unknown, which in turns reveals that values isn't of type Unknown on line 450. So I've just exposed something that was wasn't typed correctly.

All test failures except one of the Ubuntu ones can be attributed to this recent change in PyTensor. I am not sure what caused the remaining one, but I'm pretty confident it is unrelated.

Edit: Rebased on main branch and added one more set of type hints.

codecov · 2024-08-28T18:25:47Z

Codecov Report

Attention: Patch coverage is 97.56098% with 1 line in your changes missing coverage. Please review.

Project coverage is 92.74%. Comparing base (938aff4) to head (8565965).
Report is 3 commits behind head on main.

Files with missing lines	Patch %	Lines
pymc/data.py	88.88%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #7473      +/-   ##
==========================================
+ Coverage   92.69%   92.74%   +0.04%     
==========================================
  Files         104      104              
  Lines       17409    17376      -33     
==========================================
- Hits        16138    16115      -23     
+ Misses       1271     1261      -10

Files with missing lines	Coverage Δ
pymc/model/core.py	`92.78% <100.00%> (+0.95%)`	⬆️
pymc/data.py	`89.22% <88.88%> (+0.13%)`	⬆️

twiecki · 2024-09-06T11:18:05Z

Thanks @thomasaarholt, there are still mypy failures, do you have a handle on them or would you like some help?

thomasaarholt · 2024-09-06T13:28:16Z

@twiecki Thanks! The remaining one is one that I'm not sure on how best to handle:
From the mypy github action:

pymc/data.py:439: error: Incompatible types in assignment (expression has type "Sequence[str | None]", variable has type "Sequence[str] | None")

This stems from the determine_coords function having a return signature for dim that is Sequence[str | None]. But the dim type hint from higher up in the def Data constructor only allows a sequence of Sequence[str]. The sequence can be of None inner type here.

Help much appreciated!

thomasaarholt · 2024-10-04T10:01:41Z

There! @twiecki wanna take a look? I went through the logic related to the above failing mypy error and think I found a good resolution. See commit description of the last two commits.

The failing test is unrelated.

twiecki · 2024-10-04T11:45:50Z

Any idea about that failing test?

================================== FAILURES ===================================
_________________ TestMarginalVsLatent.testLatentMultioutput __________________

self = <tests.gp.test_gp.TestMarginalVsLatent object at 0x0000029EDDB84820>

    def testLatentMultioutput(self):
        n_outputs = 2
        X = np.random.randn(20, 3)
        y = np.random.randn(n_outputs, 20)
        Xnew = np.random.randn(30, 3)
        pnew = np.random.randn(n_outputs, 30)
    
        with pm.Model() as latent_model:
            cov_func = pm.gp.cov.ExpQuad(3, [0.1, 0.2, 0.3])
            mean_func = pm.gp.mean.Constant(0.5)
            latent_gp = pm.gp.Latent(mean_func=mean_func, cov_func=cov_func)
            latent_f = latent_gp.prior("f", X, n_outputs=n_outputs, reparameterize=True)
            latent_p = latent_gp.conditional("p", Xnew)
    
        with pm.Model() as marginal_model:
            cov_func = pm.gp.cov.ExpQuad(3, [0.1, 0.2, 0.3])
            mean_func = pm.gp.mean.Constant(0.5)
            marginal_gp = pm.gp.Marginal(mean_func=mean_func, cov_func=cov_func)
            marginal_f = marginal_gp.marginal_likelihood("f", X, y, sigma=0.0)
            marginal_p = marginal_gp.conditional("p", Xnew)
    
        assert tuple(latent_f.shape.eval()) == tuple(marginal_f.shape.eval()) == y.shape
        assert tuple(latent_p.shape.eval()) == tuple(marginal_p.shape.eval()) == pnew.shape
    
        chol = np.linalg.cholesky(cov_func(X).eval())
        v = np.linalg.solve(chol, (y - 0.5).T)
        A = np.linalg.solve(chol, cov_func(X, Xnew).eval()).T
        mu_cond = mean_func(Xnew).eval() + (A @ v).T
        cov_cond = cov_func(Xnew, Xnew).eval() - A @ A.T
    
        with pm.Model() as numpy_model:
            numpy_p = pm.MvNormal.dist(mu=pt.as_tensor(mu_cond), cov=pt.as_tensor(cov_cond))
    
        latent_rv_logp = pm.logp(latent_p, pnew)
        marginal_rv_logp = pm.logp(marginal_p, pnew)
        numpy_rv_logp = pm.logp(numpy_p, pnew)
    
        assert (
            latent_rv_logp.shape.eval()
            == marginal_rv_logp.shape.eval()
            == numpy_rv_logp.shape.eval()
        )
    
>       npt.assert_allclose(latent_rv_logp.eval(), marginal_rv_logp.eval(), atol=5)

tests\gp\test_gp.py:412: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

args = (<function assert_allclose.<locals>.compare at 0x0000029EE75B9AB0>, array([-41.92778875, -45.52605[201](https://github.com/pymc-devs/pymc/actions/runs/11177316239/job/31072575657?pr=7473#step:7:202)]), array([-40.07459855, -65.84149075]))
kwds = {'equal_nan': True, 'err_msg': '', 'header': 'Not equal to tolerance rtol=1e-07, atol=5', 'verbose': True}

    @wraps(func)
    def inner(*args, **kwds):
        with self._recreate_cm():
>           return func(*args, **kwds)
E           AssertionError: 
E           Not equal to tolerance rtol=1e-07, atol=5
E           
E           Mismatched elements: 1 / 2 (50%)
E           Max absolute difference: 20.31543874
E           Max relative difference: 0.30855071
E            x: array([-41.927789, -45.526052])
E            y: array([-40.074599, -65.841491])

C:\Miniconda3\envs\pymc-test\lib\contextlib.py:79: AssertionError

thomasaarholt · 2024-10-04T12:26:04Z

I assume that if we reran it (maybe you can trigger that?) it would pass. Seeing as it fails only on windows and not Linux/mac, I assumed it was a result of variance due to random sampling.

ricardoV94 · 2024-10-04T13:45:02Z

I restarted it

thomasaarholt · 2024-10-04T16:06:18Z

Ran passing everything. Ready for review!

pymc/model/core.py

ricardoV94 · 2024-10-09T07:48:01Z

@thomasaarholt the PR looks good, but since it touches on such as fundamental functionality I've asked for further reviews.

You have some conflicts that need to be resolved as well.

Thanks for the work so far!

lucianopaz

Nice work @thomasaarholt. The original code was very hard to read, and yours is much better now. I think that your refactor could be extended to a few other places. I left comments in the relevant parts.

pymc/data.py

pymc/model/core.py

thomasaarholt · 2024-10-09T15:34:29Z

@lucianopaz you tagged the wrong person ;)

Please see my reply about the get_context classmethod. Happy to get rid of more stuff (there are a few things more that are unused in model/core.py), but I'll prefer to do that in an other PR.

lucianopaz · 2024-10-10T08:30:54Z

@lucianopaz you tagged the wrong person ;)

Oops, how embarrassing! Sorry about that 😬

Please see my reply about the get_context classmethod. Happy to get rid of more stuff (there are a few things more that are unused in model/core.py), but I'll prefer to do that in an other PR.

Everything looks fine now. Thanks for the great work!

lucianopaz · 2024-10-10T08:34:52Z

@ricardoV94, I remember seeing the same test error on one of your earlier PRs. Was it fixed now? @thomasaarholt, maybe if you rebase on top of the latest main, the test will pass.

get_context returns an instance of a Model, not a ContextMeta object We don't need the typevar, since we don't use it for anything special

All of these are supported on python>=3.9.

We create a global instance of it within this module, which is similar to how it worked before, where a `context_class` attribute was attached to the Model class. We inherit from threading.local to ensure thread safety when working with models on multiple threads. See pymc-devs#1552 for the reasoning. This is already tested in `test_thread_safety`.

UNSET is the instance of the _UnsetType type. We should be typing the latter here.

We use the new ModelManager.parent_context property to reliably set any parent context, or else set it to None.

We set this directly on the class as a classmethod, which is clearer than going via the metaclass.

The original function does not behave as I expected. In the following example I expected that it would return only the final model, not root. This method is not used anywhere in the pymc codebase, so I have dropped it from the codebase. I originally included the following code to replace it, but since it is not used anyway, it is better to remove it. ```python` @classmethod def get_contexts(cls) -> list[Model]: """Return a list of the currently active model contexts.""" return MODEL_MANAGER.active_contexts ``` Example for testing behaviour in current main branch: ```python import pymc as pm with pm.Model(name="root") as root: print([c.name for c in pm.Model.get_contexts()]) with pm.Model(name="first") as first: print([c.name for c in pm.Model.get_contexts()]) with pm.Model(name="m_with_model_None", model=None) as m_with_model_None: # This one doesn't make much sense: print([c.name for c in pm.Model.get_contexts()]) ```

We only keep the __call__ method, which is necessary to keep the model context itself active during that model's __init__.

@ricardoV94

In pymc/distributions/distribution.py, this change allows the type checker to infer that `rv_out` can only be a TensorVariable. Thanks to @ricardoV94 for type hint on rv_var.

I originally tried numpy's ArrayLike, replacing Sequence entirely, but then I realized that ArrayLike also allows non-sequences like integers and floats. I am not certain if `values="a string"` should be legal. With the type hint sequence, it is. Might be more accurate, but verbose to use `list | tuple | set | np.ndarray | None`.

…unction We don't want to allow the user to pass a `dims=[None, None]` to our function, but current behaviour set `dims=[None] * N` at the end of `determine_coords`. To handle this, I created a `new_dims` with a larger type scope which matches the return type of `dims` in `determine_coords`. Then I did the same within def Data to support this new type hint.

The only case where dims=[None, ...] is when the user has passed dims=None. Since the user passed dims=None, they shouldn't be expecting any coords to match that dimension. Thus we don't need to try to add any more coords to the model.

@ricardoV94

* Type get_context correctly get_context returns an instance of a Model, not a ContextMeta object We don't need the typevar, since we don't use it for anything special * Import from future to use delayed evaluation of annotations All of these are supported on python>=3.9. * New ModelManager class for managing model contexts We create a global instance of it within this module, which is similar to how it worked before, where a `context_class` attribute was attached to the Model class. We inherit from threading.local to ensure thread safety when working with models on multiple threads. See pymc-devs#1552 for the reasoning. This is already tested in `test_thread_safety`. * Model class is now the context manager directly * Fix type of UNSET in type definition UNSET is the instance of the _UnsetType type. We should be typing the latter here. * Set model parent in init rather than in __new__ We use the new ModelManager.parent_context property to reliably set any parent context, or else set it to None. * Replace get_context in metaclass with classmethod We set this directly on the class as a classmethod, which is clearer than going via the metaclass. * Remove get_contexts from metaclass The original function does not behave as I expected. In the following example I expected that it would return only the final model, not root. This method is not used anywhere in the pymc codebase, so I have dropped it from the codebase. I originally included the following code to replace it, but since it is not used anyway, it is better to remove it. ```python` @classmethod def get_contexts(cls) -> list[Model]: """Return a list of the currently active model contexts.""" return MODEL_MANAGER.active_contexts ``` Example for testing behaviour in current main branch: ```python import pymc as pm with pm.Model(name="root") as root: print([c.name for c in pm.Model.get_contexts()]) with pm.Model(name="first") as first: print([c.name for c in pm.Model.get_contexts()]) with pm.Model(name="m_with_model_None", model=None) as m_with_model_None: # This one doesn't make much sense: print([c.name for c in pm.Model.get_contexts()]) ``` * Simplify ContextMeta We only keep the __call__ method, which is necessary to keep the model context itself active during that model's __init__. * Type Model.register_rv for for downstream typing In pymc/distributions/distribution.py, this change allows the type checker to infer that `rv_out` can only be a TensorVariable. Thanks to @ricardoV94 for type hint on rv_var. * Include np.ndarray as possible type for coord values I originally tried numpy's ArrayLike, replacing Sequence entirely, but then I realized that ArrayLike also allows non-sequences like integers and floats. I am not certain if `values="a string"` should be legal. With the type hint sequence, it is. Might be more accurate, but verbose to use `list | tuple | set | np.ndarray | None`. * Use function-scoped new_dims to handle type hint varying throughout function We don't want to allow the user to pass a `dims=[None, None]` to our function, but current behaviour set `dims=[None] * N` at the end of `determine_coords`. To handle this, I created a `new_dims` with a larger type scope which matches the return type of `dims` in `determine_coords`. Then I did the same within def Data to support this new type hint. * Fix case of dims = [None, None, ...] The only case where dims=[None, ...] is when the user has passed dims=None. Since the user passed dims=None, they shouldn't be expecting any coords to match that dimension. Thus we don't need to try to add any more coords to the model. * Remove unused hack

thomasaarholt force-pushed the thomasaarholt/modelclass branch from d203e0f to 5e08d1e Compare August 28, 2024 17:54

thomasaarholt force-pushed the thomasaarholt/modelclass branch from 5e08d1e to 383d08e Compare September 4, 2024 13:18

thomasaarholt force-pushed the thomasaarholt/modelclass branch from 7e1243e to 74665c8 Compare October 4, 2024 09:15

twiecki requested a review from ricardoV94 October 4, 2024 11:46

ricardoV94 reviewed Oct 9, 2024

View reviewed changes

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

ricardoV94 requested a review from lucianopaz October 9, 2024 07:47

ricardoV94 added the maintenance label Oct 9, 2024

thomasaarholt force-pushed the thomasaarholt/modelclass branch 2 times, most recently from b6c74de to 95bcf19 Compare October 9, 2024 09:40

lucianopaz reviewed Oct 9, 2024

View reviewed changes

pymc/data.py Show resolved Hide resolved

pymc/data.py Show resolved Hide resolved

pymc/model/core.py Show resolved Hide resolved

pymc/model/core.py Show resolved Hide resolved

thomasaarholt force-pushed the thomasaarholt/modelclass branch from c367caa to 95bcf19 Compare October 9, 2024 18:57

lucianopaz approved these changes Oct 10, 2024

View reviewed changes

thomasaarholt added 5 commits October 10, 2024 11:06

Type get_context correctly

ddf6f10

get_context returns an instance of a Model, not a ContextMeta object We don't need the typevar, since we don't use it for anything special

Import from future to use delayed evaluation of annotations

74f6616

All of these are supported on python>=3.9.

Model class is now the context manager directly

d5beee6

Fix type of UNSET in type definition

e2cefc3

UNSET is the instance of the _UnsetType type. We should be typing the latter here.

thomasaarholt added 8 commits October 10, 2024 11:06

Set model parent in init rather than in __new__

5a1f0b8

We use the new ModelManager.parent_context property to reliably set any parent context, or else set it to None.

Replace get_context in metaclass with classmethod

8846db9

We set this directly on the class as a classmethod, which is clearer than going via the metaclass.

Simplify ContextMeta

560cd3b

We only keep the __call__ method, which is necessary to keep the model context itself active during that model's __init__.

Type Model.register_rv for for downstream typing

6ff3d5d

In pymc/distributions/distribution.py, this change allows the type checker to infer that `rv_out` can only be a TensorVariable. Thanks to @ricardoV94 for type hint on rv_var.

Fix case of dims = [None, None, ...]

dbbb9a2

The only case where dims=[None, ...] is when the user has passed dims=None. Since the user passed dims=None, they shouldn't be expecting any coords to match that dimension. Thus we don't need to try to add any more coords to the model.

thomasaarholt force-pushed the thomasaarholt/modelclass branch from 95bcf19 to dbbb9a2 Compare October 10, 2024 09:07

Remove unused hack

8565965

ricardoV94 merged commit 81cfe3c into pymc-devs:main Oct 10, 2024
22 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Simplify Model new and metaclass #7473

Simplify Model new and metaclass #7473

thomasaarholt commented Aug 23, 2024 •

edited

Loading

thomasaarholt commented Aug 23, 2024 •

edited

Loading

codecov bot commented Aug 28, 2024 •

edited

Loading

twiecki commented Sep 6, 2024

thomasaarholt commented Sep 6, 2024

thomasaarholt commented Oct 4, 2024

twiecki commented Oct 4, 2024

thomasaarholt commented Oct 4, 2024

ricardoV94 commented Oct 4, 2024

thomasaarholt commented Oct 4, 2024

ricardoV94 commented Oct 9, 2024

lucianopaz left a comment •

edited

Loading

thomasaarholt commented Oct 9, 2024

lucianopaz commented Oct 10, 2024

lucianopaz commented Oct 10, 2024

Simplify Model __new__ and metaclass #7473

Simplify Model __new__ and metaclass #7473

Conversation

thomasaarholt commented Aug 23, 2024 • edited Loading

Description

Checklist

Type of change

thomasaarholt commented Aug 23, 2024 • edited Loading

codecov bot commented Aug 28, 2024 • edited Loading

Codecov Report

twiecki commented Sep 6, 2024

thomasaarholt commented Sep 6, 2024

thomasaarholt commented Oct 4, 2024

twiecki commented Oct 4, 2024

thomasaarholt commented Oct 4, 2024

ricardoV94 commented Oct 4, 2024

thomasaarholt commented Oct 4, 2024

ricardoV94 commented Oct 9, 2024

lucianopaz left a comment • edited Loading

Choose a reason for hiding this comment

thomasaarholt commented Oct 9, 2024

lucianopaz commented Oct 10, 2024

lucianopaz commented Oct 10, 2024

Simplify Model new and metaclass #7473

Simplify Model new and metaclass #7473

thomasaarholt commented Aug 23, 2024 •

edited

Loading

thomasaarholt commented Aug 23, 2024 •

edited

Loading

codecov bot commented Aug 28, 2024 •

edited

Loading

lucianopaz left a comment •

edited

Loading