How to make sense of the use of `yield` in PyMC models?

I am not a user of PyMC myself, but recently I stumbled upon this article that showed a snippet of some PyMC model:

def linear_regression(x):
    scale = yield tfd.HalfCauchy(0, 1)
    coefs = yield tfd.Normal(tf.zeros(x.shape[1]), 1, )
    predictions = yield tfd.Normal(tf.linalg.matvec(x, coefs), scale)
    return predictions

The author suggested that users

would be uncomfortable with bar = yield foo

Uncomfortable indeed I am. I tried to make sense of this generator, but couldn't see how it can be used.

This is my thought process. If I do foo = linear_regression(bar) and execute foo (e.g. next(foo)), it will return the value of scale to me. However, this will also turn the local variable scale to None. Similarly, if foo is executed again, I can get the value of coefs, but the local coefs would become None. With both local scale and coefs being None, how can predictions be evaluated?

Or is there a way to evaluate foo without triggering the yield on scale and coefs, and directly yield on predictions?

What is the black magic here? Help needed.

Solution

Disclosure: I'm the author of the original linked article.

I think your main misunderstanding is this: Python generators can not only yield values to you, but you can also send back values to generators using generator.send(). Thus, bar = yield foo would yield foo to you; the generator will wait until you send it another value (which can be None, which is what happens if you just call next(generator)!), assign that value to bar, and then continue running the generator.

Here's a simple example:

>>> def add_one_generator():
...     x = 0
...     while True:
...         x = yield x + 1
...
>>> gen = add_one_generator()
>>> y = gen.send(None)  # First sent value must be None, to start the generator
>>> print(y)
1
>>> z = gen.send(2)
>>> print(z)
3

Notice that when I send(2), the generator assigns the sent value to x, and then resumes execution. In this case, that just means yield x + 1 again, which is why the yielded z is 3.

For more info on this pattern and why it might be useful, take a look at this StackOverflow answer.

Here's some pseudocode that brings us closer to how things will (probably) work in PyMC4:

>>> def linear_regression(x):
...     scale = yield tfd.HalfCauchy(0, 1)
...     coefs = yield tfd.Normal(tf.zeros(x.shape[1]), 1, )
...     predictions = yield tfd.Normal(tf.linalg.matvec(x, coefs), scale)
...     return predictions
>>> model = linear_regression(data)
>>> next_distribution = model.send(None)
>>> scale = pymc_do_things(next_distribution)
>>> coefs = pymc_do_things(model.send(scale))
>>> predictions = pymc_do_things(model.send(coefs))