Fix pymc3 to work with latest theano-pymc master

[kyleabeauchamp]

The current pymc master branch is pinned to the latest theano release, but there have been some breaking changes since that release (e.g., Random and jax_dispatch changes). This PR fixes these compatibility issues. Note this is not expected to pass initially until we sort out the version bump for theano.

[twiecki]

There is a Theano-PyMC 1.0.13 which you can try to pin.

[kyleabeauchamp]

Several of the remaining test failures seem like a possible regression in theano or pymc3 somewhere. The following one pops up in several locations, and there are a few other pickle-related failures as well:

    def test_model_roundtrip(self):
        m = self.model
        for proto in range(pickle.HIGHEST_PROTOCOL + 1):
            try:
>               s = pickle.dumps(m, proto)
E               AttributeError: Can't pickle local object 'add_compile_configvars.<locals>.<lambda>'

[twiecki]

CC @michaelosthege @brandonwillard

[twiecki]

@kyleabeauchamp Were you able to test with aesara-devs/aesara#242 ?

[kyleabeauchamp]

I haven't tested, sorry. I tried to edit requirements.txt and trigger CI/CD but I don't think that works because we're using an environment install path that doesn't support direct git access.

[twiecki]

@kyleabeauchamp Maybe you can just run the failing test locally?

[kyleabeauchamp]

I can confirm that pytest -v pymc3/tests/test_pickling.py is now working with the theano-pymc PR, so that PR does fix some things. The remaining tests are running now.

[kyleabeauchamp]

Looking at local tests, there are some other tests failing as well. You can see them in the CI/CD logs too if you look.

self = <theano.compile.ops.Rebroadcast object at 0x7faa084385b0>, x = normal_rv.out

    def make_node(self, x):
        if self.axis.keys() and (x.ndim <= max(self.axis.keys())):
>           raise ValueError("Trying to rebroadcast non-existent dimension")
E           ValueError: Trying to rebroadcast non-existent dimension

Theano-PyMC/theano/compile/ops.py:719: ValueError
[...]
FAILED pymc3/tests/test_sampling.py::TestSample::test_sample_init - ValueError: Trying to rebroadcast non-existent dimension
FAILED pymc3/tests/test_sampling.py::test_exec_nuts_init[advi] - ValueError: Trying to rebroadcast non-existent dimension
FAILED pymc3/tests/test_variational_inference.py::test_replacements_in_sample_node_aevb[NormalizingFlowGroup: {'flow': 'scale'}] - ValueError: Trying to rebroadcast non-existent dimension
[...]

[twiecki]

/pre-commit-run

[michaelosthege]

I'm looking into aesara-devs/aesara#243

It would be great if either of you can already deal with the DeprecationWarning: Function change_flags is now deprecated! Use theano.config.change_flags instead! on this branch.

[michaelosthege]

Looks like the switch from theano.sandbox.rng_mrg.MRG_RandomStreams to theano.tensor.random.utils.RandomStream blew it up.
The two interfaces seem to be not completely identical. I'm not sure on which side it should get fixed. Waiting for input from @brandonwillard.

[brandonwillard]

Some—if not all—of the current errors are due to pymc3.variational.opvi's use of Theano's old broken RNG interface.

Specifically, lines like this wrongly assume that the first argument to an RNG sampler is always the output size/shape. This is—and was—incorrect according to the NumPy interface Theano's RNGs were supposed to emulate, and the new interface fixed that, so it absolutely won't work from here on.

After changing those lines to specify size=shape, pymc3.tests.test_sampling.TestSample.test_sample_init passes.

Also, these commits need to be squashed (e.g. all the RNG changes in one commit and the JAX ones in another).

[kyleabeauchamp]

It seems like there are still pickle failures in the CI/CD logs, even after we bumped to theano 1.0.14. E.g.,

>       dump = pickle.dumps(three_var_aevb_approx)
E       TypeError: cannot pickle 'module' object

[twiecki]

I wonder if the problem is now in the VI module? Is there a way to see which "module"?

[brandonwillard]

Multiple errors are due to the variational tests not having consistent float[32|64] types, and at least one other appears to be caused by a numerical error falling outside of a fixed range.

The former can likely be fixed by correctly setting the test input types, and the latter may be due to a new RNG seeding caused by the change to RandomStream.

[brandonwillard]

I've reorganized the commits and made some minor updates to the variational tests, so you'll need to rebase locally.

Once aesara-devs/aesara#252 goes through we can cut a release for it and pin that version in this PR. Otherwise, most of the test failures should be fixed after that, but I think there's still one remaining.

[brandonwillard]

It looks like there are some unexpected NullType gradients in pymc3.tests.test_variational_inference.test_fit_oo under the SVGD-mini parameter values. This problem isn't shared by MRG_RandomStream, so I've changed RandomStream back to MRG_RandomStream.

The reason MRG_RandomStream isn't producing NullTypes in this case might be due to how it represents uniform(a, b). Instead of representing uniform(a, b) as a single Op, MRG_RandomStream constructs a graph of uniform(a, b) via uniform(0, 1) * (b - a) + a, where uniform(0, 1) is the underlying RNG Op.

[codecov]

Codecov Report

Merging #4382 (a3cb024) into master (91993d8) will increase coverage by 0.00%.
The diff coverage is 72.72%.

@@           Coverage Diff           @@
##           master    #4382   +/-   ##
=======================================
  Coverage   88.07%   88.08%           
=======================================
  Files          88       88           
  Lines       14476    14475    -1     
=======================================
  Hits        12750    12750           
+ Misses       1726     1725    -1     

Impacted Files	Coverage Δ
pymc3/sampling_jax.py	0.00% <0.00%> (ø)
pymc3/tests/helpers.py	61.66% <100.00%> (ø)
pymc3/theanof.py	90.56% <100.00%> (ø)
pymc3/variational/opvi.py	86.51% <100.00%> (ø)

[twiecki]

Thanks @kyleabeauchamp and @brandonwillard, really glad to have PyMC3 work again with recent Theano.