final-api.py

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


linear_regression.model_generator(tf.zeros([3, 10]))  # Shape is irrelevant here

# Out[8]:
# <generator object _pm_compiled_model_generator at 0x107a5c5c8>