#@title Licensed under the Apache License, Version 2.0 (the "License"); { display-mode: "form" }
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Explanation: Auto-Batched Joint Distributions: A Gentle Tutorial
Copyright 2020 The TensorFlow Authors.
Licensed under the Apache License, Version 2.0 (the "License");
End of explanation
"""


#@title Import and set ups{ display-mode: "form" }

import functools
import numpy as np

import tensorflow.compat.v2 as tf
tf.enable_v2_behavior()

import tensorflow_probability as tfp

tfd = tfp.distributions

"""
Explanation: <table class="tfo-notebook-buttons" align="left">
  <td>
    <a target="_blank" href="https://www.tensorflow.org/probability/examples/Modeling_with_JointDistribution"><img src="https://www.tensorflow.org/images/tf_logo_32px.png" />View on TensorFlow.org</a>
  </td>
  <td>
    <a target="_blank" href="https://colab.research.google.com/github/tensorflow/probability/blob/main/tensorflow_probability/examples/jupyter_notebooks/Modeling_with_JointDistribution.ipynb"><img src="https://www.tensorflow.org/images/colab_logo_32px.png" />Run in Google Colab</a>
  </td>
  <td>
    <a target="_blank" href="https://github.com/tensorflow/probability/blob/main/tensorflow_probability/examples/jupyter_notebooks/Modeling_with_JointDistribution.ipynb"><img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />View source on GitHub</a>
  </td>
  <td>
    <a href="https://storage.googleapis.com/tensorflow_docs/probability/tensorflow_probability/examples/jupyter_notebooks/Modeling_with_JointDistribution.ipynb"><img src="https://www.tensorflow.org/images/download_logo_32px.png" />Download notebook</a>
  </td>
</table>

Introduction
TensorFlow Probability (TFP) offers a number of JointDistribution abstractions that make probabilistic inference easier by allowing a user to easily express a probabilistic graphical model in a near-mathematical form; the abstraction generates methods for sampling from the model and evaluating the log probability of samples from the model. In this tutorial, we review "autobatched" variants, which were developed after the original JointDistribution abstractions. Relative to the original, non-autobatched abstractions, the autobatched versions are simpler to use and more ergonomic, allowing many models to be expressed with less boilerplate. In this colab, we explore a simple model in (perhaps tedious) detail, making clear the problems autobatching solves, and (hopefully) teaching the reader more about TFP shape concepts along the way.
Prior to the introduction of autobatching, there were a few different variants of JointDistribution, corresponding to different syntactic styles for expressing probabilistic models: JointDistributionSequential, JointDistributionNamed, andJointDistributionCoroutine. Auobatching exists as a mixin, so we now have AutoBatched variants of all of these. In this tutorial, we explore the differences between JointDistributionSequential and JointDistributionSequentialAutoBatched; however, everything we do here is applicable to the other variants with essentially no changes.
Dependencies & Prerequisites
End of explanation
"""


X = np.arange(7)
X

"""
Explanation: Prerequisite: A Bayesian Regression Problem
We'll consider a very simple Bayesian regression scenario:
$$
\begin{align}
m & \sim \text{Normal}(0, 1) \
b & \sim \text{Normal}(0, 1) \
Y & \sim \text{Normal}(mX + b, 1)
\end{align}
$$
In this model, m and b are drawn from standard normals, and the observations Y are drawn from a normal distribution whose mean depends on the random variables m and b, and some (nonrandom, known) covariates X. (For simplicity, in this example, we assume the scale of all random variables is known.)
To perform inference in this model, we'd need to know both the covariates X and the observations Y, but for the purposes of this tutorial, we'll only need X, so we define a simple dummy X:
End of explanation
"""


jds = tfd.JointDistributionSequential([
    tfd.Normal(loc=0., scale=1.),   # m
    tfd.Normal(loc=0., scale=1.),   # b
    lambda b, m: tfd.Normal(loc=m*X + b, scale=1.) # Y
])

"""
Explanation: Desiderata
In probabilistic inference, we often want to perform two basic operations:
- sample: Drawing samples from the model.
- log_prob: Computing the log probability of a sample from the model.
The key contribution of TFP's JointDistribution abstractions (as well as of many other approaches to probabilistic programming) is to allow users to write a model once and have access to both sample and log_prob computations.
Noting that we have 7 points in our data set (X.shape = (7,)), we can now state the desiderata for an excellent JointDistribution:

sample() should produce a list of Tensors having shape [(), (), (7,)], corresponding to the scalar slope, scalar bias, and vector observations, respectively.
log_prob(sample()) should produce a scalar: the log probability of a particular slope, bias, and observations.
sample([5, 3]) should produce a list of Tensors having shape [(5, 3), (5, 3), (5, 3, 7)], representing a (5, 3)-batch of samples from the model.
log_prob(sample([5, 3])) should produce a Tensor with shape (5, 3).

We'll now look at a succession of JointDistribution models, see how to achieve the above desiderata, and hopefully learn a little more about TFP shapes along the way. 
Spoiler alert: The approach that satisfies the above  desiderata without added boilerplate is autobatching. 
First Attempt; JointDistributionSequential
End of explanation
"""


dists, sample = jds.sample_distributions()
sample

"""
Explanation: This is more or less a direct translation of the model into code. The slope m and bias b are straightforward. Y is defined using a lambda-function: the general pattern is that a lambda-function of $k$ arguments in a JointDistributionSequential (JDS) uses the previous $k$ distributions in the model. Note the "reverse" order.
We'll call sample_distributions, which returns both a sample and the underlying "sub-distributions" that were used to generate the sample. (We could have produced just the sample by calling sample; later in the tutorial it will be convenient to have the distributions as well.) The sample we produce is fine:
End of explanation
"""


jds.log_prob(sample)

"""
Explanation: But log_prob produces a result with an undesired shape:
End of explanation
"""


try:
  jds.sample([5, 3])
except tf.errors.InvalidArgumentError as e:
  print(e)

"""
Explanation: And multiple sampling doesn't work:
End of explanation
"""


sample

"""
Explanation: Let's try to understand what's going wrong.
A Brief Review: Batch and Event Shape
In TFP, an ordinary (not a JointDistribution) probability distribution has an event shape and a batch shape, and understanding the difference is crucial to effective use of TFP:

Event shape describes the shape of a single draw from the distribution; the draw may be dependent across dimensions. For scalar distributions, the event shape is []. For a 5-dimensional MultivariateNormal, the event shape is [5].
Batch shape describes independent, not identically distributed draws, aka a "batch" of distributions. Representing a batch of distributions in a single Python object is one of the key ways TFP achieves efficiency at scale.

For our purposes, a critical fact to keep in mind is that if we call log_prob on a single sample from a distribution, the result will always have a shape that matches (i.e., has as rightmost dimensions) the batch shape.
For a more in-depth discussion of shapes, see the "Understanding TensorFlow Distributions Shapes" tutorial.
Why Doesn't log_prob(sample()) Produce a Scalar?
Let's use our knowledge of batch and event shape to explore what's happening with log_prob(sample()). Here's our sample again:
End of explanation
"""


dists

"""
Explanation: And here are our distributions:
End of explanation
"""


log_prob_parts = [dist.log_prob(s) for (dist, s) in zip(dists, sample)]
log_prob_parts

np.sum(log_prob_parts) - jds.log_prob(sample)

"""
Explanation: The log probability is computed by summing the log probabilities of the sub-distributions at the (matched) elements of the parts:
End of explanation
"""


dists[2]

"""
Explanation: So, one level of explanation is that the log probability calculation is returning a 7-Tensor because the third subcomponent of log_prob_parts is a 7-Tensor. But why?
Well, we see that the last element of dists, which corresponds to our distribution over Y in the mathematial formulation, has a batch_shape of [7]. In other words, our distribution over Y is a batch of 7 independent normals (with different means and, in this case, the same scale).
We now understand what's wrong: in JDS, the distribution over Y has batch_shape=[7], a sample from the JDS represents scalars for m and b and a "batch" of 7 independent normals. and log_prob computes 7 separate log-probabilities, each of which represents the log probability of drawing m and b and a single observation Y[i] at some X[i].
Fixing log_prob(sample()) with Independent
Recall that dists[2] has event_shape=[] and batch_shape=[7]:
End of explanation
"""


y_dist_i = tfd.Independent(dists[2], reinterpreted_batch_ndims=1)
y_dist_i

"""
Explanation: By using TFP's Independent metadistribution, which converts batch dimensions to event dimensions, we can convert this into a distribution with event_shape=[7] and batch_shape=[] (we'll rename it y_dist_i because it's a distribution on Y, with the _i standing in for our Independent wrapping):
End of explanation
"""


y_dist_i.log_prob(sample[2])

"""
Explanation: Now, the log_prob of a 7-vector is a scalar:
End of explanation
"""


y_dist_i.log_prob(sample[2]) - tf.reduce_sum(dists[2].log_prob(sample[2]))

"""
Explanation: Under the covers, Independent sums over the batch:
End of explanation
"""


jds_i = tfd.JointDistributionSequential([
    tfd.Normal(loc=0., scale=1.),   # m
    tfd.Normal(loc=0., scale=1.),   # b
    lambda b, m: tfd.Independent(   # Y
        tfd.Normal(loc=m*X + b, scale=1.),
        reinterpreted_batch_ndims=1)
])

jds_i.log_prob(sample)

"""
Explanation: And indeed, we can use this to construct a new jds_i (the i again stands for Independent) where log_prob returns a scalar:
End of explanation
"""


try:
  jds_i.sample([5, 3])
except tf.errors.InvalidArgumentError as e:
  print(e)

"""
Explanation: A couple notes:
- jds_i.log_prob(s) is not the same as tf.reduce_sum(jds.log_prob(s)). The former produces the "correct" log probability of the joint distribution. The latter sums over a 7-Tensor, each element of which is the sum of the log probability of m, b, and a single element of the log probability of Y, so it overcounts m and b. (log_prob(m) + log_prob(b) + log_prob(Y) returns a result rather than throwing an exception because TFP follows TF and NumPy's broadcasting rules; adding a scalar to a vector produces a vector-sized result.)
- In this particular case, we could have solved the problem and achieved the same result using MultivariateNormalDiag instead of Independent(Normal(...)). MultivariateNormalDiag is a vector-valued distribution (i.e., it already has vector event-shape). Indeeed MultivariateNormalDiag could be (but isn't) implemented as a composition of Independent and Normal. It's worthwhile to remember that given a vector V, samples from n1 = Normal(loc=V), and n2 = MultivariateNormalDiag(loc=V) are indistinguishable; the difference beween these distributions is that n1.log_prob(n1.sample()) is a vector and n2.log_prob(n2.sample()) is a scalar.
Multiple Samples?
Drawing multiple samples still doesn't work:
End of explanation
"""


m = tfd.Normal(0., 1.).sample([5, 3])
try:
  m * X
except tf.errors.InvalidArgumentError as e:
  print(e)

"""
Explanation: Let's think about why. When we call jds_i.sample([5, 3]), we'll first draw samples for m and b, each with shape (5, 3). Next, we're going to try to construct a Normal distribution via:
tfd.Normal(loc=m*X + b, scale=1.)
But if m has shape (5, 3) and X has shape 7, we can't multiply them together, and indeed this is the error we're hitting:
End of explanation
"""


m[..., tf.newaxis].shape

(m[..., tf.newaxis] * X).shape

"""
Explanation: To resolve this issue, let's think about what properties the distribution over Y has to have. If we've called jds_i.sample([5, 3]), then we know m and b will both have shape (5, 3). What shape should a call to sample on the Y distribution produce? The obvious answer is (5, 3, 7): for each batch point, we want a sample with the same size as X. We can achieve this by using TensorFlow's broadcasting capabilities, adding extra dimensions:
End of explanation
"""


jds_ia = tfd.JointDistributionSequential([
    tfd.Normal(loc=0., scale=1.),   # m
    tfd.Normal(loc=0., scale=1.),   # b
    lambda b, m: tfd.Independent(   # Y
        tfd.Normal(loc=m[..., tf.newaxis]*X + b[..., tf.newaxis], scale=1.),
        reinterpreted_batch_ndims=1)
])

shaped_sample = jds_ia.sample([5, 3])
shaped_sample

jds_ia.log_prob(shaped_sample)

"""
Explanation: Adding an axis to both m and b, we can define a new JDS that supports multiple samples:
End of explanation
"""


(jds_ia.log_prob(shaped_sample)[3, 1] -
 jds_i.log_prob([shaped_sample[0][3, 1],
                 shaped_sample[1][3, 1],
                 shaped_sample[2][3, 1, :]]))

"""
Explanation: As an extra check, we'll verify that the log probability for a single batch point matches what we had before:
End of explanation
"""


jds_ab = tfd.JointDistributionSequentialAutoBatched([
    tfd.Normal(loc=0., scale=1.),   # m
    tfd.Normal(loc=0., scale=1.),   # b
    lambda b, m: tfd.Normal(loc=m*X + b, scale=1.) # Y
])

jds_ab.log_prob(jds.sample())

shaped_sample = jds_ab.sample([5, 3])
jds_ab.log_prob(shaped_sample)

jds_ab.log_prob(shaped_sample) - jds_ia.log_prob(shaped_sample)

"""
Explanation: <a id='AutoBatching-For-The-Win'></a>
AutoBatching For The Win
Excellent! We now have a version of JointDistribution that handles all our desiderata: log_prob returns a scalar thanks to the use of tfd.Independent, and multiple samples work now that we fixed broadcasting by adding extra axes.
What if I told you there was an easier, better way? There is, and it's called JointDistributionSequentialAutoBatched (JDSAB):
End of explanation
"""


jds.batch_shape

jds_i.batch_shape

jds_ia.batch_shape

jds_ab.batch_shape

"""
Explanation: How does this work? While you could attempt to read the code for a deep understanding, we'll give a brief overview which is sufficient for most use cases:
- Recall that our first problem was that our distribution for Y had batch_shape=[7] and event_shape=[], and we used Independent to convert the batch dimension to an event dimension. JDSAB ignores the batch shapes of component distributions; instead it treats batch shape as an overall property of the model, which is assumed to be [] (unless specified otherwise by setting batch_ndims &gt; 0). The effect is equivalent to using tfd.Independent to convert all batch dimensions of component distributions into event dimensions, as we did manually above.
- Our second problem was a need to massage the shapes of m and b so that they could broadcast appropriately with X when creating multiple samples. With JDSAB, you write a model to generate a single sample, and we "lift" the entire model to generate multiple samples using TensorFlow's vectorized_map. (This feature is analagous to JAX's vmap.)
Exploring the batch shape issue in more detail, we can compare the batch shapes of our original "bad" joint distribution jds, our batch-fixed distributions jds_i and jds_ia, and our autobatched jds_ab:
End of explanation
"""


X = np.arange(14).reshape((2, 7))
X

"""
Explanation: We see that the original jds has subdistributions with different batch shapes. jds_i and jds_ia fix this by creating subdistributions with the same (empty) batch shape. jds_ab has only a single (empty) batch shape.
It's worth noting that JointDistributionSequentialAutoBatched offers some additional generality for free. Suppose we make the covariates X (and, implicitly, the observations Y) two-dimensional:
End of explanation
"""


jds_ab = tfd.JointDistributionSequentialAutoBatched([
    tfd.Normal(loc=0., scale=1.),   # m
    tfd.Normal(loc=0., scale=1.),   # b
    lambda b, m: tfd.Normal(loc=m*X + b, scale=1.) # Y
])

shaped_sample = jds_ab.sample([5, 3])
shaped_sample

jds_ab.log_prob(shaped_sample)

"""
Explanation: Our JointDistributionSequentialAutoBatched works with no changes (we need to redefine the model because the shape of X is cached by jds_ab.log_prob):
End of explanation
"""


jds_ia = tfd.JointDistributionSequential([
    tfd.Normal(loc=0., scale=1.),   # m
    tfd.Normal(loc=0., scale=1.),   # b
    lambda b, m: tfd.Independent(   # Y
        tfd.Normal(loc=m[..., tf.newaxis]*X + b[..., tf.newaxis], scale=1.),
        reinterpreted_batch_ndims=1)
])

try:
  jds_ia.sample([5, 3])
except tf.errors.InvalidArgumentError as e:
  print(e)

"""
Explanation: On the other hand, our carefully crafted JointDistributionSequential no longer works:
End of explanation
"""


