Require all step methods to return stats #6313

michaelosthege · 2022-11-18T21:02:39Z

The reason for these changes are the resulting simplification of code, including simpler branching and less type ambiguity.

This change may break custom step methods, but whoever had enough skill to implement a custom step method should have no trouble to append a , [] to the return of their (a)step method.

Closes #6270

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

Reuires all step methods to return stats from their step/asteo method.
The BlockedStep.generates_stats attribute was removed.

Bugfixes / New features

Removed a branch in Compound.step that supported namedtuple stats which are in violation of the BlockedStep.step signature.

Docs / Maintenance

Added & fixed some type hints on step and astep methods. (arraystep.py now passing mypy!)

codecov · 2022-11-18T21:39:09Z

Codecov Report

Merging #6313 (da7e14a) into main (4acd98e) will increase coverage by 0.05%.
The diff coverage is 88.98%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #6313      +/-   ##
==========================================
+ Coverage   94.17%   94.22%   +0.05%     
==========================================
  Files         111      111              
  Lines       23908    23865      -43     
==========================================
- Hits        22515    22487      -28     
+ Misses       1393     1378      -15

Impacted Files	Coverage Δ
pymc/step_methods/hmc/hmc.py	`92.85% <ø> (-0.13%)`	⬇️
pymc/step_methods/hmc/nuts.py	`97.29% <ø> (-0.02%)`	⬇️
pymc/step_methods/metropolis.py	`83.76% <76.78%> (+0.39%)`	⬆️
pymc/aesaraf.py	`94.86% <100.00%> (ø)`
pymc/blocking.py	`95.55% <100.00%> (+0.20%)`	⬆️
pymc/sampling/mcmc.py	`92.21% <100.00%> (-0.06%)`	⬇️
pymc/sampling/parallel.py	`87.76% <100.00%> (+0.43%)`	⬆️
pymc/sampling/population.py	`72.22% <100.00%> (+0.79%)`	⬆️
pymc/step_methods/arraystep.py	`94.73% <100.00%> (+0.41%)`	⬆️
pymc/step_methods/compound.py	`100.00% <100.00%> (+15.38%)`	⬆️
... and 4 more

The reason for this change is the resulting simplification of code, including simpler branching and less type ambiguity. At the same time it allowed for fixing of a lot of type hints and method signatures on step methods. Closes pymc-devs#6270

michaelosthege · 2022-11-18T23:55:13Z

pymc/step_methods/arraystep.py

-    def astep(
-        self, apoint: RaveledVars, point: PointType, *args
-    ) -> Union[RaveledVars, Tuple[RaveledVars, StatsType]]:
+    def astep(self, apoint: RaveledVars, *args) -> Tuple[RaveledVars, StatsType]:


The exact composition of the args is specified via the constructor parameters 🤐

michaelosthege · 2022-11-18T23:57:45Z

pymc/step_methods/metropolis.py

-    def astep(self, q0: RaveledVars, logp) -> Tuple[RaveledVars, List[Dict[str, Any]]]:
-
+    def astep(self, q0: RaveledVars, *args) -> Tuple[RaveledVars, StatsType]:
+        logp = args[0]


BinaryMetropolis is not even covered by our tests!

michaelosthege · 2022-11-18T23:59:11Z

pymc/step_methods/metropolis.py

+    def astep(self, q0: RaveledVars, *args) -> Tuple[RaveledVars, StatsType]:
+        logp: Callable[[RaveledVars], np.ndarray] = args[0]


BinaryGibbsMetropolis is also not covered by tests!

michaelosthege · 2022-11-19T00:58:07Z

pymc/step_methods/metropolis.py

-        if rng is None:
-            rng = nr
-        return rng.normal(scale=self.s)
+        return (rng or nr).normal(scale=self.s)


These changes resolve two mypy errors:

rng is no longer reassigned, and

the (rng or nr) has type Generator | Module("numpy.random") which resolves Item "None" of "Optional[Generator]" has no attribute "normal"

Looks much less readable though

ricardoV94 · 2022-11-19T10:24:21Z

We shouldn't remove the trust_input. It allows evaluating the Aesara function much faster because it doesn't check the input types are valid/need to be converted every time it's called

michaelosthege · 2022-11-19T11:38:40Z

We shouldn't remove the trust_input. It allows evaluating the Aesara function much faster because it doesn't check the input types are valid/need to be converted every time it's called

But we're not using the trust_input anywhere in the codebase! Or is it a feature of the Aesara function that is not correctly advertised by the API?
Because mypy complained about it

michaelosthege · 2022-11-19T11:41:51Z

Ah I see, it's some feature of the function... I'll investigate and drop that commit.

Anything else you want to get changed, @ricardoV94 ?

pymc/step_methods/metropolis.py

michaelosthege self-assigned this Nov 18, 2022

michaelosthege added maintenance major Include in major changes release notes section labels Nov 18, 2022

michaelosthege added this to the v4.4.0 milestone Nov 18, 2022

michaelosthege force-pushed the issue-6270 branch from 3bc8c36 to f21ef72 Compare November 18, 2022 21:19

michaelosthege marked this pull request as ready for review November 18, 2022 23:49

Require all step methods to return stats

35e7c59

The reason for this change is the resulting simplification of code, including simpler branching and less type ambiguity. At the same time it allowed for fixing of a lot of type hints and method signatures on step methods. Closes pymc-devs#6270

michaelosthege force-pushed the issue-6270 branch from f21ef72 to 9545b62 Compare November 19, 2022 00:51

michaelosthege commented Nov 19, 2022

View reviewed changes

michaelosthege requested review from ricardoV94 and aseyboldt November 19, 2022 00:59

ricardoV94 reviewed Nov 19, 2022

View reviewed changes

pymc/step_methods/metropolis.py Outdated Show resolved Hide resolved

michaelosthege force-pushed the issue-6270 branch from 9545b62 to 9e09aa4 Compare November 19, 2022 12:35

Fix remaining type problems in metropolis.py

da7e14a

michaelosthege force-pushed the issue-6270 branch from 9e09aa4 to da7e14a Compare November 19, 2022 12:40

ricardoV94 approved these changes Nov 19, 2022

View reviewed changes

michaelosthege merged commit 2c3f544 into pymc-devs:main Nov 19, 2022

michaelosthege deleted the issue-6270 branch November 19, 2022 13:24

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Require all step methods to return stats #6313

Require all step methods to return stats #6313

michaelosthege commented Nov 18, 2022 •

edited

Loading

codecov bot commented Nov 18, 2022 •

edited

Loading

michaelosthege Nov 18, 2022

michaelosthege Nov 18, 2022

michaelosthege Nov 18, 2022

michaelosthege Nov 19, 2022

ricardoV94 Nov 19, 2022

ricardoV94 commented Nov 19, 2022

michaelosthege commented Nov 19, 2022 •

edited

Loading

michaelosthege commented Nov 19, 2022

		def astep(self, q0: RaveledVars, *args) -> Tuple[RaveledVars, StatsType]:
		logp: Callable[[RaveledVars], np.ndarray] = args[0]

Require all step methods to return stats #6313

Require all step methods to return stats #6313

Conversation

michaelosthege commented Nov 18, 2022 • edited Loading

Major / Breaking Changes

Bugfixes / New features

Docs / Maintenance

codecov bot commented Nov 18, 2022 • edited Loading

Codecov Report

michaelosthege Nov 18, 2022

Choose a reason for hiding this comment

michaelosthege Nov 18, 2022

Choose a reason for hiding this comment

michaelosthege Nov 18, 2022

Choose a reason for hiding this comment

michaelosthege Nov 19, 2022

Choose a reason for hiding this comment

ricardoV94 Nov 19, 2022

Choose a reason for hiding this comment

ricardoV94 commented Nov 19, 2022

michaelosthege commented Nov 19, 2022 • edited Loading

michaelosthege commented Nov 19, 2022

michaelosthege commented Nov 18, 2022 •

edited

Loading

codecov bot commented Nov 18, 2022 •

edited

Loading

michaelosthege commented Nov 19, 2022 •

edited

Loading