# Copyright 2018 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://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.
# ============================================================================
"""MaskedAutoregressiveFlow bijector."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

# Dependency imports
import numpy as np
import six
import tensorflow as tf

from tensorflow_probability.python.bijectors import bijector
from tensorflow_probability.python.internal import dtype_util
from tensorflow_probability.python.internal import tensorshape_util
from tensorflow_probability.python.math.numeric import clip_by_value_preserve_gradient

__all__ = [
    "AutoregressiveLayer",
    "MaskedAutoregressiveFlow",
    "masked_autoregressive_default_template",
    "masked_dense",
]

class MaskedAutoregressiveFlow(bijector.Bijector):
  """Affine MaskedAutoregressiveFlow bijector.

  The affine autoregressive flow [(Papamakarios et al., 2016)][3] provides a
  relatively simple framework for user-specified (deep) architectures to learn a
  distribution over continuous events. Regarding terminology,

    "Autoregressive models decompose the joint density as a product of
    conditionals, and model each conditional in turn. Normalizing flows
    transform a base density (e.g. a standard Gaussian) into the target density
    by an invertible transformation with tractable Jacobian."
    [(Papamakarios et al., 2016)][3]

  In other words, the "autoregressive property" is equivalent to the
  decomposition, `p(x) = prod{ p(x[perm[i]] | x[perm[0:i]]) : i=0, ..., d }`
  where `perm` is some permutation of `{0, ..., d}`. In the simple case where
  the permutation is identity this reduces to:
  `p(x) = prod{ p(x[i] | x[0:i]) : i=0, ..., d }`. The provided
  `shift_and_log_scale_fn`, `masked_autoregressive_default_template`, achieves
  this property by zeroing out weights in its `masked_dense` layers.

  In TensorFlow Probability, "normalizing flows" are implemented as
  `tfp.bijectors.Bijector`s. The `forward` "autoregression" is implemented
  using a `tf.while_loop` and a deep neural network (DNN) with masked weights
  such that the autoregressive property is automatically met in the `inverse`.

  A `TransformedDistribution` using `MaskedAutoregressiveFlow(...)` uses the
  (expensive) forward-mode calculation to draw samples and the (cheap)
  reverse-mode calculation to compute log-probabilities. Conversely, a
  `TransformedDistribution` using `Invert(MaskedAutoregressiveFlow(...))` uses
  the (expensive) forward-mode calculation to compute log-probabilities and the
  (cheap) reverse-mode calculation to compute samples.  See "Example Use"
  [below] for more details.

  Given a `shift_and_log_scale_fn`, the forward and inverse transformations are
  (a sequence of) affine transformations. A "valid" `shift_and_log_scale_fn`
  must compute each `shift` (aka `loc` or "mu" in [Germain et al. (2015)][1])
  and `log(scale)` (aka "alpha" in [Germain et al. (2015)][1]) such that each
  are broadcastable with the arguments to `forward` and `inverse`, i.e., such
  that the calculations in `forward`, `inverse` [below] are possible.

  For convenience, `masked_autoregressive_default_template` is offered as a
  possible `shift_and_log_scale_fn` function. It implements the MADE
  architecture [(Germain et al., 2015)][1]. MADE is a feed-forward network that
  computes a `shift` and `log(scale)` using `masked_dense` layers in a deep
  neural network. Weights are masked to ensure the autoregressive property. It
  is possible that this architecture is suboptimal for your task. To build
  alternative networks, either change the arguments to
  `masked_autoregressive_default_template`, use the `masked_dense` function to
  roll-out your own, or use some other architecture, e.g., using `tf.layers`.

  Warning: no attempt is made to validate that the `shift_and_log_scale_fn`
  enforces the "autoregressive property".

  Assuming `shift_and_log_scale_fn` has valid shape and autoregressive
  semantics, the forward transformation is

  ```python
  def forward(x):
    y = zeros_like(x)
    event_size = x.shape[-event_dims:].num_elements()
    for _ in range(event_size):
      shift, log_scale = shift_and_log_scale_fn(y)
      y = x * tf.exp(log_scale) + shift
    return y
  ```

  and the inverse transformation is

  ```python
  def inverse(y):
    shift, log_scale = shift_and_log_scale_fn(y)
    return (y - shift) / tf.exp(log_scale)
  ```

  Notice that the `inverse` does not need a for-loop. This is because in the
  forward pass each calculation of `shift` and `log_scale` is based on the `y`
  calculated so far (not `x`). In the `inverse`, the `y` is fully known, thus is
  equivalent to the scaling used in `forward` after `event_size` passes, i.e.,
  the "last" `y` used to compute `shift`, `log_scale`. (Roughly speaking, this
  also proves the transform is bijective.)

  #### Examples

  ```python
  tfd = tfp.distributions
  tfb = tfp.bijectors

  dims = 5

  # A common choice for a normalizing flow is to use a Gaussian for the base
  # distribution. (However, any continuous distribution would work.) E.g.,
  maf = tfd.TransformedDistribution(
      distribution=tfd.Normal(loc=0., scale=1.),
      bijector=tfb.MaskedAutoregressiveFlow(
          shift_and_log_scale_fn=tfb.masked_autoregressive_default_template(
              hidden_layers=[512, 512])),
      event_shape=[dims])

  x = maf.sample()  # Expensive; uses `tf.while_loop`, no Bijector caching.
  maf.log_prob(x)   # Almost free; uses Bijector caching.
  maf.log_prob(0.)  # Cheap; no `tf.while_loop` despite no Bijector caching.

  # [Papamakarios et al. (2016)][3] also describe an Inverse Autoregressive
  # Flow [(Kingma et al., 2016)][2]:
  iaf = tfd.TransformedDistribution(
      distribution=tfd.Normal(loc=0., scale=1.),
      bijector=tfb.Invert(tfb.MaskedAutoregressiveFlow(
          shift_and_log_scale_fn=tfb.masked_autoregressive_default_template(
              hidden_layers=[512, 512]))),
      event_shape=[dims])

  x = iaf.sample()  # Cheap; no `tf.while_loop` despite no Bijector caching.
  iaf.log_prob(x)   # Almost free; uses Bijector caching.
  iaf.log_prob(0.)  # Expensive; uses `tf.while_loop`, no Bijector caching.

  # In many (if not most) cases the default `shift_and_log_scale_fn` will be a
  # poor choice. Here's an example of using a "shift only" version and with a
  # different number/depth of hidden layers.
  shift_only = True
  maf_no_scale_hidden2 = tfd.TransformedDistribution(
      distribution=tfd.Normal(loc=0., scale=1.),
      bijector=tfb.MaskedAutoregressiveFlow(
          tfb.masked_autoregressive_default_template(
              hidden_layers=[32],
              shift_only=shift_only),
          is_constant_jacobian=shift_only),
      event_shape=[dims])
  ```

  #### References

  [1]: Mathieu Germain, Karol Gregor, Iain Murray, and Hugo Larochelle. MADE:
       Masked Autoencoder for Distribution Estimation. In _International
       Conference on Machine Learning_, 2015. https://arxiv.org/abs/1502.03509

  [2]: Diederik P. Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya
       Sutskever, and Max Welling. Improving Variational Inference with Inverse
       Autoregressive Flow. In _Neural Information Processing Systems_, 2016.
       https://arxiv.org/abs/1606.04934

  [3]: George Papamakarios, Theo Pavlakou, and Iain Murray. Masked
       Autoregressive Flow for Density Estimation. In _Neural Information
       Processing Systems_, 2017. https://arxiv.org/abs/1705.07057
  """

  def __init__(self,
               shift_and_log_scale_fn,
               is_constant_jacobian=False,
               validate_args=False,
               unroll_loop=False,
               event_ndims=1,
               name=None):
    """Creates the MaskedAutoregressiveFlow bijector.

    Args:
      shift_and_log_scale_fn: Python `callable` which computes `shift` and
        `log_scale` from both the forward domain (`x`) and the inverse domain
        (`y`). Calculation must respect the "autoregressive property" (see class
        docstring). Suggested default
        `masked_autoregressive_default_template(hidden_layers=...)`.
        Typically the function contains `tf.Variables` and is wrapped using
        `tf.make_template`. Returning `None` for either (both) `shift`,
        `log_scale` is equivalent to (but more efficient than) returning zero.
      is_constant_jacobian: Python `bool`. Default: `False`. When `True` the
        implementation assumes `log_scale` does not depend on the forward domain
        (`x`) or inverse domain (`y`) values. (No validation is made;
        `is_constant_jacobian=False` is always safe but possibly computationally
        inefficient.)
      validate_args: Python `bool` indicating whether arguments should be
        checked for correctness.
      unroll_loop: Python `bool` indicating whether the `tf.while_loop` in
        `_forward` should be replaced with a static for loop. Requires that
        the final dimension of `x` be known at graph construction time. Defaults
        to `False`.
      event_ndims: Python `integer`, the intrinsic dimensionality of this
        bijector. 1 corresponds to a simple vector autoregressive bijector as
        implemented by the `masked_autoregressive_default_template`, 2 might be
        useful for a 2D convolutional `shift_and_log_scale_fn` and so on.
      name: Python `str`, name given to ops managed by this object.
    """
    name = name or "masked_autoregressive_flow"
    self._shift_and_log_scale_fn = shift_and_log_scale_fn
    self._unroll_loop = unroll_loop
    self._event_ndims = event_ndims
    super(MaskedAutoregressiveFlow, self).__init__(
        forward_min_event_ndims=self._event_ndims,
        is_constant_jacobian=is_constant_jacobian,
        validate_args=validate_args,
        name=name)

  def _forward(self, x, **kwargs):
    static_event_size = tensorshape_util.num_elements(
        tensorshape_util.with_rank_at_least(
            x.shape, self._event_ndims)[-self._event_ndims:])

    if self._unroll_loop:
      if not static_event_size:
        raise ValueError(
            "The final {} dimensions of `x` must be known at graph "
            "construction time if `unroll_loop=True`. `x.shape: {!r}`".format(
                self._event_ndims, x.shape))
      y = tf.zeros_like(x, name="y0")

      for _ in range(static_event_size):
        shift, log_scale = self._shift_and_log_scale_fn(y, **kwargs)
        # next_y = scale * x + shift
        next_y = x
        if log_scale is not None:
          next_y *= tf.exp(log_scale)
        if shift is not None:
          next_y += shift
        y = next_y
      return y

    event_size = tf.reduce_prod(
        input_tensor=tf.shape(input=x)[-self._event_ndims:])
    y0 = tf.zeros_like(x, name="y0")
    # call the template once to ensure creation
    _ = self._shift_and_log_scale_fn(y0, **kwargs)
    def _loop_body(index, y0):
      """While-loop body for autoregression calculation."""
      # Set caching device to avoid re-getting the tf.Variable for every while
      # loop iteration.
      with tf.compat.v1.variable_scope(tf.compat.v1.get_variable_scope()) as vs:
        if vs.caching_device is None and not tf.executing_eagerly():
          vs.set_caching_device(lambda op: op.device)
        shift, log_scale = self._shift_and_log_scale_fn(y0, **kwargs)
      y = x
      if log_scale is not None:
        y *= tf.exp(log_scale)
      if shift is not None:
        y += shift
      return index + 1, y
    # If the event size is available at graph construction time, we can inform
    # the graph compiler of the maximum number of steps. If not,
    # static_event_size will be None, and the maximum_iterations argument will
    # have no effect.
    _, y = tf.while_loop(
        cond=lambda index, _: index < event_size,
        body=_loop_body,
        loop_vars=(0, y0),
        maximum_iterations=static_event_size)
    return y

  def _inverse(self, y, **kwargs):
    shift, log_scale = self._shift_and_log_scale_fn(y, **kwargs)
    x = y
    if shift is not None:
      x -= shift
    if log_scale is not None:
      x *= tf.exp(-log_scale)
    return x

  def _inverse_log_det_jacobian(self, y, **kwargs):
    _, log_scale = self._shift_and_log_scale_fn(y, **kwargs)
    if log_scale is None:
      return tf.constant(0., dtype=y.dtype, name="ildj")
    return -tf.reduce_sum(
        input_tensor=log_scale, axis=tf.range(-self._event_ndims, 0))

MASK_INCLUSIVE = "inclusive"
MASK_EXCLUSIVE = "exclusive"

def _gen_slices(num_blocks, n_in, n_out, mask_type=MASK_EXCLUSIVE):
  """Generate the slices for building an autoregressive mask."""
  # TODO(b/67594795): Better support of dynamic shape.
  slices = []
  col = 0
  d_in = n_in // num_blocks
  d_out = n_out // num_blocks
  row = d_out if mask_type == MASK_EXCLUSIVE else 0
  for _ in range(num_blocks):
    row_slice = slice(row, None)
    col_slice = slice(col, col + d_in)
    slices.append([row_slice, col_slice])
    col += d_in
    row += d_out
  return slices

def _gen_mask(num_blocks,
              n_in,
              n_out,
              mask_type=MASK_EXCLUSIVE,
              dtype=tf.float32):
  """Generate the mask for building an autoregressive dense layer."""
  # TODO(b/67594795): Better support of dynamic shape.
  mask = np.zeros([n_out, n_in], dtype=dtype.as_numpy_dtype())
  slices = _gen_slices(num_blocks, n_in, n_out, mask_type=mask_type)
  for [row_slice, col_slice] in slices:
    mask[row_slice, col_slice] = 1
  return mask

def masked_dense(inputs,
                 units,
                 num_blocks=None,
                 exclusive=False,
                 kernel_initializer=None,
                 reuse=None,
                 name=None,
                 *args,  # pylint: disable=keyword-arg-before-vararg
                 **kwargs):
  """A autoregressively masked dense layer. Analogous to `tf.layers.dense`.

  See [Germain et al. (2015)][1] for detailed explanation.

  Arguments:
    inputs: Tensor input.
    units: Python `int` scalar representing the dimensionality of the output
      space.
    num_blocks: Python `int` scalar representing the number of blocks for the
      MADE masks.
    exclusive: Python `bool` scalar representing whether to zero the diagonal of
      the mask, used for the first layer of a MADE.
    kernel_initializer: Initializer function for the weight matrix.
      If `None` (default), weights are initialized using the
      `tf.glorot_random_initializer`.
    reuse: Python `bool` scalar representing whether to reuse the weights of a
      previous layer by the same name.
    name: Python `str` used to describe ops managed by this function.
    *args: `tf.layers.dense` arguments.
    **kwargs: `tf.layers.dense` keyword arguments.

  Returns:
    Output tensor.

  Raises:
    NotImplementedError: if rightmost dimension of `inputs` is unknown prior to
      graph execution.

  #### References

  [1]: Mathieu Germain, Karol Gregor, Iain Murray, and Hugo Larochelle. MADE:
       Masked Autoencoder for Distribution Estimation. In _International
       Conference on Machine Learning_, 2015. https://arxiv.org/abs/1502.03509
  """
  # TODO(b/67594795): Better support of dynamic shape.
  input_depth = tf.compat.dimension_value(
      tensorshape_util.with_rank_at_least(inputs.shape, 1)[-1])
  if input_depth is None:
    raise NotImplementedError(
        "Rightmost dimension must be known prior to graph execution.")

  mask = _gen_mask(num_blocks, input_depth, units,
                   MASK_EXCLUSIVE if exclusive else MASK_INCLUSIVE).T

  if kernel_initializer is None:
    kernel_initializer = tf.compat.v1.glorot_normal_initializer()

  def masked_initializer(shape, dtype=None, partition_info=None):
    return mask * kernel_initializer(shape, dtype, partition_info)

  with tf.compat.v2.name_scope(name or "masked_dense"):
    layer = tf.compat.v1.layers.Dense(
        units,
        kernel_initializer=masked_initializer,
        kernel_constraint=lambda x: mask * x,
        name=name,
        dtype=dtype_util.base_dtype(inputs.dtype),
        _scope=name,
        _reuse=reuse,
        *args,  # pylint: disable=keyword-arg-before-vararg
        **kwargs)
    return layer.apply(inputs)

def masked_autoregressive_default_template(hidden_layers,
                                           shift_only=False,
                                           activation=tf.nn.relu,
                                           log_scale_min_clip=-5.,
                                           log_scale_max_clip=3.,
                                           log_scale_clip_gradient=False,
                                           name=None,
                                           *args,  # pylint: disable=keyword-arg-before-vararg
                                           **kwargs):
  """Build the Masked Autoregressive Density Estimator (Germain et al., 2015).

  This will be wrapped in a make_template to ensure the variables are only
  created once. It takes the input and returns the `loc` ("mu" in [Germain et
  al. (2015)][1]) and `log_scale` ("alpha" in [Germain et al. (2015)][1]) from
  the MADE network.

  Warning: This function uses `masked_dense` to create randomly initialized
  `tf.Variables`. It is presumed that these will be fit, just as you would any
  other neural architecture which uses `tf.layers.dense`.

  #### About Hidden Layers

  Each element of `hidden_layers` should be greater than the `input_depth`
  (i.e., `input_depth = tf.shape(input)[-1]` where `input` is the input to the
  neural network). This is necessary to ensure the autoregressivity property.

  #### About Clipping

  This function also optionally clips the `log_scale` (but possibly not its
  gradient). This is useful because if `log_scale` is too small/large it might
  underflow/overflow making it impossible for the `MaskedAutoregressiveFlow`
  bijector to implement a bijection. Additionally, the `log_scale_clip_gradient`
  `bool` indicates whether the gradient should also be clipped. The default does
  not clip the gradient; this is useful because it still provides gradient
  information (for fitting) yet solves the numerical stability problem. I.e.,
  `log_scale_clip_gradient = False` means
  `grad[exp(clip(x))] = grad[x] exp(clip(x))` rather than the usual
  `grad[clip(x)] exp(clip(x))`.

  Args:
    hidden_layers: Python `list`-like of non-negative integer, scalars
      indicating the number of units in each hidden layer. Default: `[512, 512].
    shift_only: Python `bool` indicating if only the `shift` term shall be
      computed. Default: `False`.
    activation: Activation function (callable). Explicitly setting to `None`
      implies a linear activation.
    log_scale_min_clip: `float`-like scalar `Tensor`, or a `Tensor` with the
      same shape as `log_scale`. The minimum value to clip by. Default: -5.
    log_scale_max_clip: `float`-like scalar `Tensor`, or a `Tensor` with the
      same shape as `log_scale`. The maximum value to clip by. Default: 3.
    log_scale_clip_gradient: Python `bool` indicating that the gradient of
      `tf.clip_by_value` should be preserved. Default: `False`.
    name: A name for ops managed by this function. Default:
      "masked_autoregressive_default_template".
    *args: `tf.layers.dense` arguments.
    **kwargs: `tf.layers.dense` keyword arguments.

  Returns:
    shift: `Float`-like `Tensor` of shift terms (the "mu" in
      [Germain et al.  (2015)][1]).
    log_scale: `Float`-like `Tensor` of log(scale) terms (the "alpha" in
      [Germain et al. (2015)][1]).

  Raises:
    NotImplementedError: if rightmost dimension of `inputs` is unknown prior to
      graph execution.

  #### References

  [1]: Mathieu Germain, Karol Gregor, Iain Murray, and Hugo Larochelle. MADE:
       Masked Autoencoder for Distribution Estimation. In _International
       Conference on Machine Learning_, 2015. https://arxiv.org/abs/1502.03509
  """
  name = name or "masked_autoregressive_default_template"
  with tf.compat.v2.name_scope(name):

    def _fn(x):
      """MADE parameterized via `masked_autoregressive_default_template`."""
      # TODO(b/67594795): Better support of dynamic shape.
      input_depth = tf.compat.dimension_value(
          tensorshape_util.with_rank_at_least(x.shape, 1)[-1])
      if input_depth is None:
        raise NotImplementedError(
            "Rightmost dimension must be known prior to graph execution.")
      input_shape = (
          np.int32(tensorshape_util.as_list(x.shape))
          if tensorshape_util.is_fully_defined(x.shape) else tf.shape(input=x))
      if tensorshape_util.rank(x.shape) == 1:
        x = x[tf.newaxis, ...]
      for i, units in enumerate(hidden_layers):
        x = masked_dense(
            inputs=x,
            units=units,
            num_blocks=input_depth,
            exclusive=True if i == 0 else False,
            activation=activation,
            *args,  # pylint: disable=keyword-arg-before-vararg
            **kwargs)
      x = masked_dense(
          inputs=x,
          units=(1 if shift_only else 2) * input_depth,
          num_blocks=input_depth,
          activation=None,
          *args,  # pylint: disable=keyword-arg-before-vararg
          **kwargs)
      if shift_only:
        x = tf.reshape(x, shape=input_shape)
        return x, None
      x = tf.reshape(x, shape=tf.concat([input_shape, [2]], axis=0))
      shift, log_scale = tf.unstack(x, num=2, axis=-1)
      which_clip = (
          tf.clip_by_value
          if log_scale_clip_gradient else clip_by_value_preserve_gradient)
      log_scale = which_clip(log_scale, log_scale_min_clip, log_scale_max_clip)
      return shift, log_scale

    return tf.compat.v1.make_template(name, _fn)

class AutoregressiveLayer(tf.keras.layers.Layer):
  r"""Masked Autoencoder for Distribution Estimation [Germain et al. (2015)][1].

  A `AutoregressiveLayer` takes as input a Tensor of shape `[..., event_size]`
  and returns a Tensor of shape `[..., event_size, params]`.

  The output satisfies the autoregressive property.  That is, the layer is
  configured with some permutation `ord` of `{0, ..., event_size-1}` (i.e., an
  ordering of the input dimensions), and the output `output[batch_idx, i, ...]`
  for input dimension `i` depends only on inputs `x[batch_idx, j]` where
  `ord(j) < ord(i)`.  The autoregressive property allows us to use
  `output[batch_idx, i]` to parameterize conditional distributions:
    `p(x[batch_idx, i] | x[batch_idx, ] for ord(j) < ord(i))`
  which give us a tractable distribution over input `x[batch_idx]`:
    `p(x[batch_idx]) = prod_i p(x[batch_idx, ord(i)] | x[batch_idx, ord(0:i)])`

  For example, when `params` is 2, the output of the layer can parameterize
  the location and log-scale of an autoregressive Gaussian distribution.

  #### Example

  ```python
  # Generate data -- as in Figure 1 in [Papamakarios et al. (2017)][2]).
  n = 2000
  x2 = np.random.randn(n).astype(dtype=np.float32) * 2.
  x1 = np.random.randn(n).astype(dtype=np.float32) + (x2 * x2 / 4.)
  data = np.stack([x1, x2], axis=-1)

  # Density estimation with MADE.
  made = tfb.AutoregressiveLayer(params=2, hidden_units=[10, 10])

  distribution = tfd.TransformedDistribution(
      distribution=tfd.Normal(loc=0., scale=1.),
      bijector=tfb.MaskedAutoregressiveFlow(
          lambda x: tf.unstack(made(x), num=2, axis=-1)),
      event_shape=[2])

  # Construct and fit model.
  x_ = tfkl.Input(shape=(2,), dtype=tf.float32)
  log_prob_ = distribution.log_prob(x_)
  model = tfk.Model(x_, log_prob_)

  model.compile(optimizer=tf.compat.v2.optimizers.Adam(),
                loss=lambda _, log_prob: -log_prob)

  batch_size = 25
  model.fit(x=data,
            y=np.zeros((n, 0), dtype=np.float32),
            batch_size=batch_size,
            epochs=1,
            steps_per_epoch=1,  # Usually `n // batch_size`.
            shuffle=True,
            verbose=True)

  # Use the fitted distribution.
  distribution.sample((3, 1))
  distribution.log_prob(np.ones((3, 2), dtype=np.float32))
  ```

  #### Examples: Handling Rank-2+ Tensors

  `AutoregressiveLayer` can be used as a building block to achieve different
  autoregressive structures over rank-2+ tensors.  For example, suppose we want
  to build an autoregressive distribution over images with dimension `[weight,
  height, channels]` with `channels = 3`:

   1. We can parameterize a "fully autoregressive" distribution, with
      cross-channel and within-pixel autoregressivity:
      ```
          r0    g0   b0     r0    g0   b0       r0   g0    b0
          ^   ^      ^         ^   ^   ^         ^      ^   ^
          |  /  ____/           \  |  /           \____  \  |
          | /__/                 \ | /                 \__\ |
          r1    g1   b1     r1 <- g1   b1       r1   g1 <- b1
                                               ^          |
                                                \_________/
      ```

      as:
      ```python
      # Generate random images for training data.
      images = np.random.uniform(size=(100, 8, 8, 3)).astype(np.float32)
      n, width, height, channels = images.shape

      # Reshape images to achieve desired autoregressivity.
      event_shape = [height * width * channels]
      reshaped_images = tf.reshape(images, [n, event_shape])

      # Density estimatino with MADE.
      made = tfb.AutoregressiveLayer(params=2, event_shape=event_shape,
                                     hidden_units=[20, 20], activation="relu")
      distribution = tfd.TransformedDistribution(
          distribution=tfd.Normal(loc=0., scale=1.),
          bijector=tfb.MaskedAutoregressiveFlow(
              lambda x: tf.unstack(made(x), num=2, axis=-1)),
          event_shape=event_shape)

      # Construct and fit model.
      x_ = tfkl.Input(shape=event_shape, dtype=tf.float32)
      log_prob_ = distribution.log_prob(x_)
      model = tfk.Model(x_, log_prob_)

      model.compile(optimizer=tf.compat.v2.optimizers.Adam(),
                    loss=lambda _, log_prob: -log_prob)

      batch_size = 10
      model.fit(x=data,
                y=np.zeros((n, 0), dtype=np.float32),
                batch_size=batch_size,
                epochs=10,
                steps_per_epoch=n // batch_size,
                shuffle=True,
                verbose=True)

      # Use the fitted distribution.
      distribution.sample((3, 1))
      distribution.log_prob(np.ones((5, 8, 8, 3), dtype=np.float32))
      ```

   2. We can parameterize a distribution with neither cross-channel nor
      within-pixel autoregressivity:
      ```
          r0    g0   b0
          ^     ^    ^
          |     |    |
          |     |    |
          r1    g1   b1
      ```

      as:
      ```python
      # Generate fake images.
      images = np.random.choice([0, 1], size=(100, 8, 8, 3))
      n, width, height, channels = images.shape

      # Reshape images to achieve desired autoregressivity.
      reshaped_images = np.transpose(
          np.reshape(images, [n, width * height, channels]),
          axes=[0, 2, 1])

      made = tfb.AutoregressiveLayer(params=1, event_shape=[width * height],
                                     hidden_units=[20, 20], activation="relu")

      # Density estimation with MADE.
      #
      # NOTE: Parameterize an autoregressive distribution over an event_shape of
      # [channels, width * height], with univariate Bernoulli conditional
      # distributions.
      distribution = tfd.Autoregressive(
          lambda x: tfd.Independent(
              tfd.Bernoulli(logits=tf.unstack(made(x), axis=-1)[0],
                            dtype=tf.float32),
              reinterpreted_batch_ndims=2),
          sample0=tf.zeros([channels, width * height], dtype=tf.float32))

      # Construct and fit model.
      x_ = tfkl.Input(shape=(channels, width * height), dtype=tf.float32)
      log_prob_ = distribution.log_prob(x_)
      model = tfk.Model(x_, log_prob_)

      model.compile(optimizer=tf.compat.v2.optimizers.Adam(),
                    loss=lambda _, log_prob: -log_prob)

      batch_size = 10
      model.fit(x=reshaped_images,
                y=np.zeros((n, 0), dtype=np.float32),
                batch_size=batch_size,
                epochs=10,
                steps_per_epoch=n // batch_size,
                shuffle=True,
                verbose=True)

      distribution.sample(7)
      distribution.log_prob(np.ones((4, 8, 8, 3), dtype=np.float32))
      ```

      Note that one set of weights is shared for the mapping for each channel
      from image to distribution parameters -- i.e., the mapping
      `layer(reshaped_images[..., channel, :])`, where `channel` is 0, 1, or 2.

      To use separate weights for each channel, we could construct an
      `AutoregressiveLayer` and `TransformedDistribution` for each channel, and
      combine them with a `tfd.Blockwise` distribution.

  #### References

  [1]: Mathieu Germain, Karol Gregor, Iain Murray, and Hugo Larochelle. MADE:
       Masked Autoencoder for Distribution Estimation. In _International
       Conference on Machine Learning_, 2015. https://arxiv.org/abs/1502.03509

  [2]: George Papamakarios, Theo Pavlakou, Iain Murray, Masked Autoregressive
       Flow for Density Estimation.  In _Neural Information Processing Systems_,
       2017. https://arxiv.org/abs/1705.07057
  """

  def __init__(self,
               params,
               event_shape=None,
               hidden_units=None,
               input_order="left-to-right",
               hidden_degrees="equal",
               activation=None,
               use_bias=True,
               kernel_initializer="glorot_uniform",
               validate_args=False,
               **kwargs):
    """Constructs the MADE layer.

    Arguments:
      params: Python integer specifying the number of parameters to output
        per input.
      event_shape: Python `list`-like of positive integers (or a single int),
        specifying the shape of the input to this layer, which is also the
        event_shape of the distribution parameterized by this layer.  Currently
        only rank-1 shapes are supported.  That is, event_shape must be a single
        integer.  If not specified, the event shape is inferred when this layer
        is first called or built.
      hidden_units: Python `list`-like of non-negative integers, specifying
        the number of units in each hidden layer.
      input_order: Order of degrees to the input units: 'random',
        'left-to-right', 'right-to-left', or an array of an explicit order. For
        example, 'left-to-right' builds an autoregressive model:
        `p(x) = p(x1) p(x2 | x1) ... p(xD | x<D)`.  Default: 'left-to-right'.
      hidden_degrees: Method for assigning degrees to the hidden units:
        'equal', 'random'.  If 'equal', hidden units in each layer are allocated
        equally (up to a remainder term) to each degree.  Default: 'equal'.
      activation: An activation function.  See `tf.keras.layers.Dense`. Default:
        `None`.
      use_bias: Whether or not the dense layers constructed in this layer
        should have a bias term.  See `tf.keras.layers.Dense`.  Default: `True`.
      kernel_initializer: Initializer for the kernel weights matrix.  Default:
        'glorot_uniform'.
      validate_args: Python `bool`, default `False`. When `True`, layer
        parameters are checked for validity despite possibly degrading runtime
        performance. When `False` invalid inputs may silently render incorrect
        outputs.
      **kwargs: Addition keyword arguments passed to the `tf.keras.layers.Dense`
        constructed by this layer.
    """
    super(AutoregressiveLayer, self).__init__(**kwargs)

    self._params = params
    self._event_shape = _list(event_shape) if event_shape is not None else None
    self._hidden_units = hidden_units if hidden_units is not None else []
    self._input_order_param = input_order
    self._hidden_degrees = hidden_degrees
    self._activation = activation
    self._use_bias = use_bias
    self._kernel_initializer = kernel_initializer
    self._validate_args = validate_args
    self._kwargs = kwargs

    if self._event_shape is not None:
      self._event_size = self._event_shape[-1]
      self._event_ndims = len(self._event_shape)

      if self._event_ndims != 1:
        raise ValueError("Parameter `event_shape` must describe a rank-1 shape."
                         " `event_shape: {!r}`".format(event_shape))

    # To be built in `build`.
    self._input_order = None
    self._masks = None
    self._network = None

  def build(self, input_shape):
    """See tfkl.Layer.build."""
    if self._event_shape is None:
      # `event_shape` wasn't specied at __init__, so infer from `input_shape`.
      self._event_shape = [tf.compat.dimension_value(input_shape[-1])]
      self._event_size = self._event_shape[-1]
      self._event_ndims = len(self._event_shape)
      # Should we throw if input_shape has rank > 2?

    if input_shape[-1] != self._event_shape[-1]:
      raise ValueError("Invalid final dimension of `input_shape`. "
                       "Expected `{!r}`, but got `{!r}`".format(
                           self._event_shape[-1], input_shape[-1]))

    # Construct the masks.
    self._input_order = _create_input_order(
        self._event_size, self._input_order_param)
    self._masks = _create_masks(_create_degrees(
        input_size=self._event_size,
        hidden_units=self._hidden_units,
        input_order=self._input_order,
        hidden_degrees=self._hidden_degrees))

    # In the final layer, we will produce `self._params` outputs for each of the
    # `self._event_size` inputs to `AutoregressiveLayer`.  But `masks[-1]` has
    # shape `[self._hidden_units[-1], self._event_size]`.  Thus, we need to
    # expand the mask to `[hidden_units[-1], event_size * self._params]` such
    # that all units for the same input are masked identically.  In particular,
    # we tile the mask so the j-th element of `tf.unstack(output, axis=-1)` is a
    # tensor of the j-th parameter/unit for each input.
    #
    # NOTE: Other orderings of the output could be faster -- should benchmark.
    self._masks[-1] = np.reshape(
        np.tile(self._masks[-1][..., tf.newaxis], [1, 1, self._params]),
        [self._masks[-1].shape[0], self._event_size * self._params])

    self._network = tf.keras.Sequential([
        # Starting this model with an `InputLayer` ensures that Keras will build
        # and propagate our `dtype` to each layer we add.
        tf.keras.layers.InputLayer((self._event_size,), dtype=self.dtype)
    ])

    # Input-to-hidden, hidden-to-hidden, and hidden-to-output layers:
    #  [..., self._event_size] -> [..., self._hidden_units[0]].
    #  [..., self._hidden_units[k-1]] -> [..., self._hidden_units[k]].
    #  [..., self._hidden_units[-1]] -> [..., event_size * self._params].
    layer_output_sizes = self._hidden_units + [self._event_size * self._params]
    for k in range(len(self._masks)):
      self._network.add(tf.keras.layers.Dense(
          layer_output_sizes[k],
          kernel_initializer=_make_masked_initializer(
              self._masks[k], self._kernel_initializer),
          kernel_constraint=_make_masked_constraint(self._masks[k]),
          activation=self._activation if k + 1 < len(self._masks) else None,
          use_bias=self._use_bias,
          **self._kwargs))

    # Record that the layer has been built.
    super(AutoregressiveLayer, self).build(input_shape)

  def call(self, x):
    """See tfkl.Layer.call."""
    with tf.compat.v2.name_scope(self.name or "AutoregressiveLayer_call"):
      x = tf.convert_to_tensor(value=x, dtype=self.dtype, name="x")
      input_shape = tf.shape(input=x)
      # TODO(b/67594795): Better support for dynamic shapes.
      if tensorshape_util.rank(x.shape) == 1:
        x = x[tf.newaxis, ...]
      return tf.reshape(self._network(x),
                        tf.concat([input_shape, [self._params]], axis=0))

  def compute_output_shape(self, input_shape):
    """See tfkl.Layer.compute_output_shape."""
    return input_shape + (self._params,)

  @property
  def event_shape(self):
    return self._event_shape

  @property
  def params(self):
    return self._params

def _list(xs):
  """Convert the given argument to a list."""
  try:
    return list(xs)
  except TypeError:
    return [xs]

def _create_input_order(input_size, input_order="left-to-right"):
  """Returns a degree vectors for the input."""
  if isinstance(input_order, six.string_types):
    if input_order == "left-to-right":
      return np.arange(start=1, stop=input_size + 1)
    elif input_order == "right-to-left":
      return np.arange(start=input_size, stop=0, step=-1)
    elif input_order == "random":
      ret = np.arange(start=1, stop=input_size + 1)
      np.random.shuffle(ret)
      return ret
  elif np.all(np.sort(input_order) == np.arange(1, input_size + 1)):
    return np.array(input_order)

  raise ValueError("Invalid input order: '{}'.".format(input_order))

def _create_degrees(input_size,
                    hidden_units=None,
                    input_order="left-to-right",
                    hidden_degrees="equal"):
  """Returns a list of degree vectors, one for each input and hidden layer.

  A unit with degree d can only receive input from units with degree < d. Output
  units always have the same degree as their associated input unit.

  Args:
    input_size: Number of inputs.
    hidden_units: list with the number of hidden units per layer. It does not
      include the output layer. Each hidden unit size must be at least the size
      of length (otherwise autoregressivity is not possible).
    input_order: Order of degrees to the input units: 'random', 'left-to-right',
      'right-to-left', or an array of an explicit order. For example,
      'left-to-right' builds an autoregressive model
      p(x) = p(x1) p(x2 | x1) ... p(xD | x<D).
    hidden_degrees: Method for assigning degrees to the hidden units:
      'equal', 'random'.  If 'equal', hidden units in each layer are allocated
      equally (up to a remainder term) to each degree.  Default: 'equal'.

  Raises:
    ValueError: invalid input order.
    ValueError: invalid hidden degrees.
  """
  input_order = _create_input_order(input_size, input_order)
  degrees = [input_order]

  if hidden_units is None:
    hidden_units = []

  for units in hidden_units:
    if isinstance(hidden_degrees, six.string_types):
      if hidden_degrees == "random":
        # samples from: [low, high)
        degrees.append(
            np.random.randint(low=min(np.min(degrees[-1]), input_size - 1),
                              high=input_size,
                              size=units))
      elif hidden_degrees == "equal":
        min_degree = min(np.min(degrees[-1]), input_size - 1)
        degrees.append(np.maximum(
            min_degree,
            # Evenly divide the range `[1, input_size - 1]` in to `units + 1`
            # segments, and pick the boundaries between the segments as degrees.
            np.ceil(np.arange(1, units + 1)
                    * (input_size - 1) / float(units + 1)).astype(np.int32)))
    else:
      raise ValueError('Invalid hidden order: "{}".'.format(hidden_degrees))

  return degrees

def _create_masks(degrees):
  """Returns a list of binary mask matrices enforcing autoregressivity."""
  return [
      # Create input->hidden and hidden->hidden masks.
      inp[:, np.newaxis] <= out
      for inp, out in zip(degrees[:-1], degrees[1:])
  ] + [
      # Create hidden->output mask.
      degrees[-1][:, np.newaxis] < degrees[0]
  ]

def _make_masked_initializer(mask, initializer):
  """Returns a masked version of the given initializer."""
  initializer = tf.keras.initializers.get(initializer)
  def masked_initializer(shape, dtype=None, partition_info=None):
    # If no `partition_info` is given, then don't pass it to `initializer`, as
    # `initializer` may be a `tf.compat.v2.initializers.Initializer` (which
    # don't accept a `partition_info` argument).
    if partition_info is None:
      x = initializer(shape, dtype)
    else:
      x = initializer(shape, dtype, partition_info)
    return tf.cast(mask, x.dtype) * x
  return masked_initializer

def _make_masked_constraint(mask):
  def masked_constraint(x):
    x = tf.convert_to_tensor(value=x, dtype_hint=tf.float32, name="x")
    return tf.cast(mask, x.dtype) * x
  return masked_constraint


# Copyright 2018 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://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.
# ============================================================================
"""The Multinomial distribution class."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import tensorflow.compat.v2 as tf

from tensorflow_probability.python.distributions import distribution
from tensorflow_probability.python.internal import assert_util
from tensorflow_probability.python.internal import distribution_util
from tensorflow_probability.python.internal import dtype_util
from tensorflow_probability.python.internal import reparameterization
from tensorflow_probability.python.internal import tensorshape_util

__all__ = [
    "Multinomial",
]

_multinomial_sample_note = """For each batch of counts, `value = [n_0, ...
,n_{k-1}]`, `P[value]` is the probability that after sampling `self.total_count`
draws from this Multinomial distribution, the number of draws falling in class
`j` is `n_j`. Since this definition is [exchangeable](
https://en.wikipedia.org/wiki/Exchangeable_random_variables); different
sequences have the same counts so the probability includes a combinatorial
coefficient.

Note: `value` must be a non-negative tensor with dtype `self.dtype`, have no
fractional components, and such that
`tf.reduce_sum(value, -1) = self.total_count`. Its shape must be broadcastable
with `self.probs` and `self.total_count`."""

class Multinomial(distribution.Distribution):
  """Multinomial distribution.

  This Multinomial distribution is parameterized by `probs`, a (batch of)
  length-`K` `prob` (probability) vectors (`K > 1`) such that
  `tf.reduce_sum(probs, -1) = 1`, and a `total_count` number of trials, i.e.,
  the number of trials per draw from the Multinomial. It is defined over a
  (batch of) length-`K` vector `counts` such that
  `tf.reduce_sum(counts, -1) = total_count`. The Multinomial is identically the
  Binomial distribution when `K = 2`.

  #### Mathematical Details

  The Multinomial is a distribution over `K`-class counts, i.e., a length-`K`
  vector of non-negative integer `counts = n = [n_0, ..., n_{K-1}]`.

  The probability mass function (pmf) is,

  ```none
  pmf(n; pi, N) = prod_j (pi_j)**n_j / Z
  Z = (prod_j n_j!) / N!
  ```

  where:
  * `probs = pi = [pi_0, ..., pi_{K-1}]`, `pi_j > 0`, `sum_j pi_j = 1`,
  * `total_count = N`, `N` a positive integer,
  * `Z` is the normalization constant, and,
  * `N!` denotes `N` factorial.

  Distribution parameters are automatically broadcast in all functions; see
  examples for details.

  #### Pitfalls

  The number of classes, `K`, must not exceed:
  - the largest integer representable by `self.dtype`, i.e.,
    `2**(mantissa_bits+1)` (IEE754),
  - the maximum `Tensor` index, i.e., `2**31-1`.

  In other words,

  ```python
  K <= min(2**31-1, {
    tf.float16: 2**11,
    tf.float32: 2**24,
    tf.float64: 2**53 }[param.dtype])
  ```

  Note: This condition is validated only when `self.validate_args = True`.

  #### Examples

  Create a 3-class distribution, with the 3rd class is most likely to be drawn,
  using logits.

  ```python
  logits = [-50., -43, 0]
  dist = Multinomial(total_count=4., logits=logits)
  ```

  Create a 3-class distribution, with the 3rd class is most likely to be drawn.

  ```python
  p = [.2, .3, .5]
  dist = Multinomial(total_count=4., probs=p)
  ```

  The distribution functions can be evaluated on counts.

  ```python
  # counts same shape as p.
  counts = [1., 0, 3]
  dist.prob(counts)  # Shape []

  # p will be broadcast to [[.2, .3, .5], [.2, .3, .5]] to match counts.
  counts = [[1., 2, 1], [2, 2, 0]]
  dist.prob(counts)  # Shape [2]

  # p will be broadcast to shape [5, 7, 3] to match counts.
  counts = [[...]]  # Shape [5, 7, 3]
  dist.prob(counts)  # Shape [5, 7]
  ```

  Create a 2-batch of 3-class distributions.

  ```python
  p = [[.1, .2, .7], [.3, .3, .4]]  # Shape [2, 3]
  dist = Multinomial(total_count=[4., 5], probs=p)

  counts = [[2., 1, 1], [3, 1, 1]]
  dist.prob(counts)  # Shape [2]

  dist.sample(5) # Shape [5, 2, 3]
  ```
  """

  def __init__(self,
               total_count,
               logits=None,
               probs=None,
               validate_args=False,
               allow_nan_stats=True,
               name="Multinomial"):
    """Initialize a batch of Multinomial distributions.

    Args:
      total_count: Non-negative floating point tensor with shape broadcastable
        to `[N1,..., Nm]` with `m >= 0`. Defines this as a batch of
        `N1 x ... x Nm` different Multinomial distributions. Its components
        should be equal to integer values.
      logits: Floating point tensor representing unnormalized log-probabilities
        of a positive event with shape broadcastable to
        `[N1,..., Nm, K]` `m >= 0`, and the same dtype as `total_count`. Defines
        this as a batch of `N1 x ... x Nm` different `K` class Multinomial
        distributions. Only one of `logits` or `probs` should be passed in.
      probs: Positive floating point tensor with shape broadcastable to
        `[N1,..., Nm, K]` `m >= 0` and same dtype as `total_count`. Defines
        this as a batch of `N1 x ... x Nm` different `K` class Multinomial
        distributions. `probs`'s components in the last portion of its shape
        should sum to `1`. Only one of `logits` or `probs` should be passed in.
      validate_args: Python `bool`, default `False`. When `True` distribution
        parameters are checked for validity despite possibly degrading runtime
        performance. When `False` invalid inputs may silently render incorrect
        outputs.
      allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
        (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
        result is undefined. When `False`, an exception is raised if one or
        more of the statistic's batch members are undefined.
      name: Python `str` name prefixed to Ops created by this class.
    """
    parameters = dict(locals())
    with tf.name_scope(name) as name:
      dtype = dtype_util.common_dtype([total_count, logits, probs], tf.float32)
      self._total_count = tf.convert_to_tensor(
          value=total_count, name="total_count", dtype=dtype)
      if validate_args:
        self._total_count = (
            distribution_util.embed_check_nonnegative_integer_form(
                self._total_count))
      self._logits, self._probs = distribution_util.get_logits_and_probs(
          logits=logits,
          probs=probs,
          multidimensional=True,
          validate_args=validate_args,
          name=name,
          dtype=dtype)
      self._mean_val = self._total_count[..., tf.newaxis] * self._probs
    super(Multinomial, self).__init__(
        dtype=dtype,
        reparameterization_type=reparameterization.NOT_REPARAMETERIZED,
        validate_args=validate_args,
        allow_nan_stats=allow_nan_stats,
        parameters=parameters,
        graph_parents=[self._total_count, self._logits, self._probs],
        name=name)

  @classmethod
  def _params_event_ndims(cls):
    return dict(total_count=0, logits=1, probs=1)

  @property
  def total_count(self):
    """Number of trials used to construct a sample."""
    return self._total_count

  @property
  def logits(self):
    """Vector of coordinatewise logits."""
    return self._logits

  @property
  def probs(self):
    """Probability of drawing a `1` in that coordinate."""
    return self._probs

  def _batch_shape_tensor(self):
    return tf.shape(input=self._mean_val)[:-1]

  def _batch_shape(self):
    return tensorshape_util.with_rank_at_least(self._mean_val.shape, 1)[:-1]

  def _event_shape_tensor(self):
    return tf.shape(input=self._mean_val)[-1:]

  def _event_shape(self):
    return tensorshape_util.with_rank_at_least(self._mean_val.shape, 1)[-1:]

  def _sample_n(self, n, seed=None):
    n_draws = tf.cast(self.total_count, dtype=tf.int32)
    k = self.event_shape_tensor()[0]

    return draw_sample(n, k, self.logits, n_draws, self.dtype, seed)

  @distribution_util.AppendDocstring(_multinomial_sample_note)
  def _log_prob(self, counts):
    return self._log_unnormalized_prob(counts) - self._log_normalization(counts)

  def _log_unnormalized_prob(self, counts):
    counts = self._maybe_assert_valid_sample(counts)
    return tf.reduce_sum(
        input_tensor=counts * tf.nn.log_softmax(self.logits), axis=-1)

  def _log_normalization(self, counts):
    counts = self._maybe_assert_valid_sample(counts)
    return -distribution_util.log_combinations(self.total_count, counts)

  def _mean(self):
    return tf.identity(self._mean_val)

  def _covariance(self):
    p = self.probs * tf.ones_like(
        self.total_count)[..., tf.newaxis]
    return tf.linalg.set_diag(
        -tf.matmul(self._mean_val[..., tf.newaxis],
                   p[..., tf.newaxis, :]),  # outer product
        self._variance())

  def _variance(self):
    p = self.probs * tf.ones_like(
        self.total_count)[..., tf.newaxis]
    return self._mean_val - self._mean_val * p

  def _maybe_assert_valid_sample(self, counts):
    """Check counts for proper shape, values, then return tensor version."""
    if not self.validate_args:
      return counts
    counts = distribution_util.embed_check_nonnegative_integer_form(counts)
    return distribution_util.with_dependencies([
        assert_util.assert_equal(
            self.total_count,
            tf.reduce_sum(input_tensor=counts, axis=-1),
            message="counts must sum to `self.total_count`"),
    ], counts)

def draw_sample(num_samples, num_classes, logits, num_trials, dtype, seed):
  """Sample a multinomial.

  The batch shape is given by broadcasting num_trials with
  remove_last_dimension(logits).

  Args:
    num_samples: Python int or singleton integer Tensor: number of multinomial
      samples to draw.
    num_classes: Python int or singleton integer Tensor: number of classes.
    logits: Floating Tensor with last dimension k, of (unnormalized) logit
      probabilities per class.
    num_trials: Tensor of number of categorical trials each multinomial consists
      of.  num_trials[..., tf.newaxis] must broadcast with logits.
    dtype: dtype at which to emit samples.
    seed: Random seed.

  Returns:
    samples: Tensor of given dtype and shape [n] + batch_shape + [k].
  """
  with tf.name_scope("multinomial.draw_sample"):
    # broadcast the num_trials and logits to same shape
    num_trials = tf.ones_like(
        logits[..., 0], dtype=num_trials.dtype) * num_trials
    logits = tf.ones_like(
        num_trials[..., tf.newaxis], dtype=logits.dtype) * logits

    # flatten the total_count and logits
    # flat_logits has shape [B1B2...Bm, num_classes]
    flat_logits = tf.reshape(logits, [-1, num_classes])
    flat_num_trials = num_samples * tf.reshape(num_trials, [-1])  # [B1B2...Bm]

    # Computes each logits and num_trials situation by map_fn.

    # Using just one batch tf.random.categorical call doesn't work because that
    # requires num_trials to be the same across all members of the batch of
    # logits.  This restriction makes sense for tf.random.categorical because
    # for it, num_trials is part of the returned shape.  However, the
    # multinomial sampler does not need that restriction, because it sums out
    # exactly that dimension.

    # One possibility would be to draw a batch categorical whose sample count is
    # max(num_trials) and mask out the excess ones.  However, if the elements of
    # num_trials vary widely, this can be wasteful of memory.

    # TODO(b/123763054, b/112152209): Revisit the possibility of writing this
    # with a batch categorical followed by batch unsorted_segment_sum, once both
    # of those work and are memory-efficient enough.
    def _sample_one_batch_member(args):
      logits, num_cat_samples = args[0], args[1]  # [K], []
      # x has shape [1, num_cat_samples = num_samples * num_trials]
      x = tf.random.categorical(
          logits[tf.newaxis, ...], num_cat_samples, seed=seed)
      x = tf.reshape(x, shape=[num_samples, -1])  # [num_samples, num_trials]
      x = tf.one_hot(
          x, depth=num_classes)  # [num_samples, num_trials, num_classes]
      x = tf.reduce_sum(input_tensor=x, axis=-2)  # [num_samples, num_classes]
      return tf.cast(x, dtype=dtype)

    x = tf.map_fn(
        _sample_one_batch_member, [flat_logits, flat_num_trials],
        dtype=dtype)  # [B1B2...Bm, num_samples, num_classes]

    # reshape the results to proper shape
    x = tf.transpose(a=x, perm=[1, 0, 2])
    final_shape = tf.concat([[num_samples],
                             tf.shape(input=num_trials), [num_classes]],
                            axis=0)
    x = tf.reshape(x, final_shape)

    return x


# Copyright 2018 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://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.
# ============================================================================
"""Regression components."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

# Dependency imports
import tensorflow as tf

from tensorflow_probability.python import bijectors as tfb
from tensorflow_probability.python import distributions as tfd

from tensorflow_probability.python.sts.structural_time_series import Parameter
from tensorflow_probability.python.sts.structural_time_series import StructuralTimeSeries

tfl = tf.linalg

def _zero_dimensional_mvndiag(dtype):
  """Build a zero-dimensional MVNDiag object."""
  dummy_mvndiag = tfd.MultivariateNormalDiag(
      scale_diag=tf.ones([0], dtype=dtype))
  dummy_mvndiag.covariance = lambda: dummy_mvndiag.variance()[..., tf.newaxis]
  return dummy_mvndiag

def _observe_timeseries_fn(timeseries):
  """Build an observation_noise_fn that observes a Tensor timeseries."""
  def observation_noise_fn(t):
    current_slice = timeseries[..., t, :]
    return tfd.MultivariateNormalDiag(
        loc=current_slice,
        scale_diag=tf.zeros_like(current_slice))
  return observation_noise_fn

class LinearRegression(StructuralTimeSeries):
  """Formal representation of a linear regression from provided covariates.

  This model defines a time series given by a linear combination of
  covariate time series provided in a design matrix:

  ```python
  observed_time_series = matmul(design_matrix, weights)
  ```

  The design matrix has shape `[num_timesteps, num_features]`. The weights
  are treated as an unknown random variable of size `[num_features]` (both
  components also support batch shape), and are integrated over using the same
  approximate inference tools as other model parameters, i.e., generally HMC or
  variational inference.

  This component does not itself include observation noise; it defines a
  deterministic distribution with mass at the point
  `matmul(design_matrix, weights)`. In practice, it should be combined with
  observation noise from another component such as `tfp.sts.Sum`, as
  demonstrated below.

  #### Examples

  Given `series1`, `series2` as `Tensors` each of shape `[num_timesteps]`
  representing covariate time series, we create a regression model that
  conditions on these covariates:

  ```python
  regression = tfp.sts.LinearRegression(
    design_matrix=tf.stack([series1, series2], axis=-1),
    weights_prior=tfd.Normal(loc=0., scale=1.))
  ```

  Here we've also demonstrated specifying a custom prior, using an informative
  `Normal(0., 1.)` prior instead of the default weakly-informative prior.

  As a more advanced application, we might use the design matrix to encode
  holiday effects. For example, suppose we are modeling data from the month of
  December. We can combine day-of-week seasonality with special effects for
  Christmas Eve (Dec 24), Christmas (Dec 25), and New Year's Eve (Dec 31),
  by constructing a design matrix with indicators for those dates.

  ```python
  holiday_indicators = np.zeros([31, 3])
  holiday_indicators[23, 0] = 1  # Christmas Eve
  holiday_indicators[24, 1] = 1  # Christmas Day
  holiday_indicators[30, 2] = 1  # New Year's Eve

  holidays = tfp.sts.LinearRegression(design_matrix=holiday_indicators,
                                      name='holidays')
  day_of_week = tfp.sts.Seasonal(num_seasons=7,
                                 observed_time_series=observed_time_series,
                                 name='day_of_week')
  model = tfp.sts.Sum(components=[holidays, seasonal],
                      observed_time_series=observed_time_series)
  ```

  Note that the `Sum` component in the above model also incorporates observation
  noise, with prior scale heuristically inferred from `observed_time_series`.

  In these examples, we've used a single design matrix, but batching is
  also supported. If the design matrix has batch shape, the default behavior
  constructs weights with matching batch shape, which will fit a separate
  regression for each design matrix. This can be overridden by passing an
  explicit weights prior with appropriate batch shape. For example, if each
  design matrix in a batch contains features with the same semantics
  (e.g., if they represent per-group or per-observation covariates), we might
  choose to share statistical strength by fitting a single weight vector that
  broadcasts across all design matrices:

  ```python
  design_matrix = get_batch_of_inputs()
  design_matrix.shape  # => concat([batch_shape, [num_timesteps, num_features]])

  # Construct a prior with batch shape `[]` and event shape `[num_features]`,
  # so that it describes a single vector of weights.
  weights_prior = tfd.Independent(
      tfd.StudentT(df=5,
                   loc=tf.zeros([num_features]),
                   scale=tf.ones([num_features])),
      reinterpreted_batch_ndims=1)
  linear_regression = LinearRegression(design_matrix=design_matrix,
                                       weights_prior=weights_prior)
  ```

  """

  def __init__(self,
               design_matrix,
               weights_prior=None,
               name=None):
    """Specify a linear regression model.

    Note: the statistical behavior of the regression is determined by
    the broadcasting behavior of the `weights` `Tensor`:

    * `weights_prior.batch_shape == []`: shares a single set of weights across
      all design matrices and observed time series. This may make sense if
      the features in each design matrix have the same semantics (e.g.,
      grouping observations by country, with per-country design matrices
      capturing the same set of national economic indicators per country).
    * `weights_prior.batch_shape == `design_matrix.batch_shape`: fits separate
      weights for each design matrix. If there are multiple observed time series
      for each design matrix, this shares statistical strength over those
      observations.
    * `weights_prior.batch_shape == `observed_time_series.batch_shape`: fits a
      separate regression for each individual time series.

    When modeling batches of time series, you should think carefully about
    which behavior makes sense, and specify `weights_prior` accordingly:
    the defaults may not do what you want!

    Args:
      design_matrix: float `Tensor` of shape `concat([batch_shape,
        [num_timesteps, num_features]])`. This may also optionally be
        an instance of `tf.linalg.LinearOperator`.
      weights_prior: `tfd.Distribution` representing a prior over the regression
        weights. Must have event shape `[num_features]` and batch shape
        broadcastable to the design matrix's `batch_shape`. Alternately,
        `event_shape` may be scalar (`[]`), in which case the prior is
        internally broadcast as `TransformedDistribution(weights_prior,
        tfb.Identity(), event_shape=[num_features],
        batch_shape=design_matrix.batch_shape)`. If `None`,
        defaults to `StudentT(df=5, loc=0., scale=10.)`, a weakly-informative
        prior loosely inspired by the [Stan prior choice recommendations](
        https://github.com/stan-dev/stan/wiki/Prior-Choice-Recommendations).
        Default value: `None`.
      name: the name of this model component.
        Default value: 'LinearRegression'.
    """
    with tf.compat.v1.name_scope(
        name, 'LinearRegression', values=[design_matrix]) as name:

      if not isinstance(design_matrix, tfl.LinearOperator):
        design_matrix = tfl.LinearOperatorFullMatrix(
            tf.convert_to_tensor(value=design_matrix, name='design_matrix'),
            name='design_matrix_linop')

      if tf.compat.dimension_value(design_matrix.shape[-1]) is not None:
        num_features = design_matrix.shape[-1]
      else:
        num_features = design_matrix.shape_tensor()[-1]

      # Default to a weakly-informative StudentT(df=5, 0., 10.) prior.
      if weights_prior is None:
        weights_prior = tfd.StudentT(
            df=5,
            loc=tf.zeros([], dtype=design_matrix.dtype),
            scale=10 * tf.ones([], dtype=design_matrix.dtype))
      # Sugar: if prior is static scalar, broadcast it to a default shape.
      if weights_prior.event_shape.ndims == 0:
        if design_matrix.batch_shape.is_fully_defined():
          design_matrix_batch_shape_ = design_matrix.batch_shape
        else:
          design_matrix_batch_shape_ = design_matrix.batch_shape_tensor()
        weights_prior = tfd.TransformedDistribution(
            weights_prior,
            bijector=tfb.Identity(),
            batch_shape=design_matrix_batch_shape_,
            event_shape=[num_features])

      tf.debugging.assert_same_float_dtype([design_matrix, weights_prior])

      self._design_matrix = design_matrix

      super(LinearRegression, self).__init__(
          parameters=[
              Parameter('weights', weights_prior, tfb.Identity()),
          ],
          latent_size=0,
          name=name)

  @property
  def design_matrix(self):
    """LinearOperator representing the design matrix."""
    return self._design_matrix

  def _make_state_space_model(self,
                              num_timesteps,
                              param_map,
                              initial_state_prior=None,
                              initial_step=0):

    weights = param_map['weights']  # shape: [B, num_features]
    predicted_timeseries = self.design_matrix.matmul(weights[..., tf.newaxis])

    dtype = self.design_matrix.dtype

    # Since this model has `latent_size=0`, the latent prior and
    # transition model are dummy objects (zero-dimensional MVNs).
    dummy_mvndiag = _zero_dimensional_mvndiag(dtype)
    if initial_state_prior is None:
      initial_state_prior = dummy_mvndiag

    return tfd.LinearGaussianStateSpaceModel(
        num_timesteps=num_timesteps,
        transition_matrix=tf.zeros([0, 0], dtype=dtype),
        transition_noise=dummy_mvndiag,
        observation_matrix=tf.zeros([1, 0], dtype=dtype),
        observation_noise=_observe_timeseries_fn(predicted_timeseries),
        initial_state_prior=initial_state_prior,
        initial_step=initial_step)

class SparseLinearRegression(StructuralTimeSeries):
  """Formal representation of a sparse linear regression.

  This model defines a time series given by a sparse linear combination of
  covariate time series provided in a design matrix:

  ```python
  observed_time_series = matmul(design_matrix, weights)
  ```

  This is identical to `tfp.sts.LinearRegression`, except that
  `SparseLinearRegression` uses a parameterization of a Horseshoe
  prior [1][2] to encode the assumption that many of the `weights` are zero,
  i.e., many of the covariate time series are irrelevant. See the mathematical
  details section below for further discussion. The prior parameterization used
  by `SparseLinearRegression` is more suitable for inference than that
  obtained by simply passing the equivalent `tfd.Horseshoe` prior to
  `LinearRegression`; when sparsity is desired, `SparseLinearRegression` will
  likely yield better results.

  This component does not itself include observation noise; it defines a
  deterministic distribution with mass at the point
  `matmul(design_matrix, weights)`. In practice, it should be combined with
  observation noise from another component such as `tfp.sts.Sum`, as
  demonstrated below.

  #### Examples

  Given `series1`, `series2` as `Tensors` each of shape `[num_timesteps]`
  representing covariate time series, we create a regression model that
  conditions on these covariates:

  ```python
  regression = tfp.sts.SparseLinearRegression(
    design_matrix=tf.stack([series1, series2], axis=-1),
    weights_prior_scale=0.1)
  ```

  The `weights_prior_scale` determines the level of sparsity; small
  scales encourage the weights to be sparse. In some cases, such as when
  the likelihood is iid Gaussian with known scale, the prior scale can be
  analytically related to the expected number of nonzero weights [2]; however,
  this is not the case in general for STS models.

  If the design matrix has batch dimensions, by default the model will create a
  matching batch of weights. For example, if `design_matrix.shape == [
  num_users, num_timesteps, num_features]`, by default the model will fit
  separate weights for each user, i.e., it will internally represent
  `weights.shape == [num_users, num_features]`. To share weights across some or
  all batch dimensions, you can manually specify the batch shape for the
  weights:

  ```python
  # design_matrix.shape == [num_users, num_timesteps, num_features]
  regression = tfp.sts.SparseLinearRegression(
    design_matrix=design_matrix,
    weights_batch_shape=[])  # weights.shape -> [num_features]
  ```

  #### Mathematical Details

  The basic horseshoe prior [1] is defined as a Cauchy-normal scale mixture:

  ```
  scales[i] ~ HalfCauchy(loc=0, scale=1)
  weights[i] ~ Normal(loc=0., scale=scales[i] * global_scale)`
  ```

  The Cauchy scale parameters puts substantial mass near zero, encouraging
  weights to be sparse, but their heavy tails allow weights far from zero to be
  estimated without excessive shrinkage. The horseshoe can be thought of as a
  continuous relaxation of a traditional 'spike-and-slab' discrete sparsity
  prior, in which the latent Cauchy scale mixes between 'spike'
  (`scales[i] ~= 0`) and 'slab' (`scales[i] >> 0`) regimes.

  Following the recommendations in [2], `SparseLinearRegression` implements
  a horseshoe with the following adaptations:

  - The Cauchy prior on `scales[i]` is represented as an InverseGamma-Normal
    compound.
  - The `global_scale` parameter is integrated out following a `Cauchy(0.,
    scale=weights_prior_scale)` hyperprior, which is also represented as an
    InverseGamma-Normal compound.
  - All compound distributions are implemented using a non-centered
    parameterization.

  The compound, non-centered representation defines the same marginal prior as
  the original horseshoe (up to integrating out the global scale),
  but allows samplers to mix more efficiently through the heavy tails; for
  variational inference, the compound representation implicity expands the
  representational power of the variational model.

  Note that we do not yet implement the regularized ('Finnish') horseshoe,
  proposed in [2] for models with weak likelihoods, because the likelihood
  in STS models is typically Gaussian, where it's not clear that additional
  regularization is appropriate. If you need this functionality, please
  email tfprobability@tensorflow.org.

  The full prior parameterization implemented in `SparseLinearRegression` is
  as follows:

  ```
  # Sample global_scale from Cauchy(0, scale=weights_prior_scale).
  global_scale_variance ~ InverseGamma(alpha=0.5, beta=0.5)
  global_scale_noncentered ~ HalfNormal(loc=0, scale=1)
  global_scale = (global_scale_noncentered *
                  sqrt(global_scale_variance) *
                  weights_prior_scale)

  # Sample local_scales from Cauchy(0, 1).
  local_scale_variances[i] ~ InverseGamma(alpha=0.5, beta=0.5)
  local_scales_noncentered[i] ~ HalfNormal(loc=0, scale=1)
  local_scales[i] = local_scales_noncentered[i] * sqrt(local_scale_variances[i])

  weights[i] ~ Normal(loc=0., scale=local_scales[i] * global_scale)
  ```

  #### References

  [1]: Carvalho, C., Polson, N. and Scott, J. Handling Sparsity via the
    Horseshoe. AISTATS (2009).
    http://proceedings.mlr.press/v5/carvalho09a/carvalho09a.pdf
  [2]: Juho Piironen, Aki Vehtari. Sparsity information and regularization in
    the horseshoe and other shrinkage priors (2017).
    https://arxiv.org/abs/1707.01694

  """

  def __init__(self,
               design_matrix,
               weights_prior_scale=0.1,
               weights_batch_shape=None,
               name=None):
    """Specify a sparse linear regression model.

    Args:
      design_matrix: float `Tensor` of shape `concat([batch_shape,
        [num_timesteps, num_features]])`. This may also optionally be
        an instance of `tf.linalg.LinearOperator`.
      weights_prior_scale: float `Tensor` defining the scale of the Horseshoe
        prior on regression weights. Small values encourage the weights to be
        sparse. The shape must broadcast with `weights_batch_shape`.
        Default value: `0.1`.
      weights_batch_shape: if `None`, defaults to
        `design_matrix.batch_shape_tensor()`. Must broadcast with the batch
        shape of `design_matrix`.
        Default value: `None`.
      name: the name of this model component.
        Default value: 'SparseLinearRegression'.
    """
    with tf.compat.v1.name_scope(
        name, 'SparseLinearRegression',
        values=[design_matrix, weights_prior_scale]) as name:

      if not isinstance(design_matrix, tfl.LinearOperator):
        design_matrix = tfl.LinearOperatorFullMatrix(
            tf.convert_to_tensor(value=design_matrix, name='design_matrix'),
            name='design_matrix_linop')

      if tf.compat.dimension_value(design_matrix.shape[-1]) is not None:
        num_features = design_matrix.shape[-1]
      else:
        num_features = design_matrix.shape_tensor()[-1]

      if weights_batch_shape is None:
        weights_batch_shape = design_matrix.batch_shape_tensor()
      else:
        weights_batch_shape = tf.convert_to_tensor(value=weights_batch_shape,
                                                   dtype=tf.int32)
      weights_shape = tf.concat([weights_batch_shape, [num_features]], axis=0)

      dtype = design_matrix.dtype

      self._design_matrix = design_matrix
      self._weights_prior_scale = weights_prior_scale

      ones_like_weights_batch = tf.ones(weights_batch_shape, dtype=dtype)
      ones_like_weights = tf.ones(weights_shape, dtype=dtype)
      super(SparseLinearRegression, self).__init__(
          parameters=[
              Parameter('global_scale_variance',
                        prior=tfd.InverseGamma(
                            0.5 * ones_like_weights_batch,
                            0.5 * ones_like_weights_batch),
                        bijector=tfb.Softplus()),
              Parameter('global_scale_noncentered',
                        prior=tfd.HalfNormal(
                            scale=ones_like_weights_batch),
                        bijector=tfb.Softplus()),
              Parameter('local_scale_variances',
                        prior=tfd.Independent(tfd.InverseGamma(
                            0.5 * ones_like_weights,
                            0.5 * ones_like_weights),
                                              reinterpreted_batch_ndims=1),
                        bijector=tfb.Softplus()),
              Parameter('local_scales_noncentered',
                        prior=tfd.Independent(tfd.HalfNormal(
                            scale=ones_like_weights),
                                              reinterpreted_batch_ndims=1),
                        bijector=tfb.Softplus()),
              Parameter('weights_noncentered',
                        prior=tfd.Independent(tfd.Normal(
                            loc=tf.zeros_like(ones_like_weights),
                            scale=ones_like_weights),
                                              reinterpreted_batch_ndims=1),
                        bijector=tfb.Identity())
          ],
          latent_size=0,
          name=name)

  @property
  def design_matrix(self):
    """LinearOperator representing the design matrix."""
    return self._design_matrix

  @property
  def weights_prior_scale(self):
    return self._weights_prior_scale

  def params_to_weights(self,
                        global_scale_variance,
                        global_scale_noncentered,
                        local_scale_variances,
                        local_scales_noncentered,
                        weights_noncentered):
    """Build regression weights from model parameters."""
    global_scale = (global_scale_noncentered *
                    tf.sqrt(global_scale_variance) *
                    self.weights_prior_scale)

    local_scales = local_scales_noncentered * tf.sqrt(local_scale_variances)
    return weights_noncentered * local_scales * global_scale[..., tf.newaxis]

  def _make_state_space_model(self,
                              num_timesteps,
                              param_map,
                              initial_state_prior=None,
                              initial_step=0):

    weights = self.params_to_weights(**param_map)
    predicted_timeseries = self.design_matrix.matmul(weights[..., tf.newaxis])

    dtype = self.design_matrix.dtype

    # Since this model has `latent_size=0`, the latent prior and
    # transition model are dummy objects (zero-dimensional MVNs).
    dummy_mvndiag = _zero_dimensional_mvndiag(dtype)
    if initial_state_prior is None:
      initial_state_prior = dummy_mvndiag

    return tfd.LinearGaussianStateSpaceModel(
        num_timesteps=num_timesteps,
        transition_matrix=tf.zeros([0, 0], dtype=dtype),
        transition_noise=dummy_mvndiag,
        observation_matrix=tf.zeros([1, 0], dtype=dtype),
        observation_noise=_observe_timeseries_fn(predicted_timeseries),
        initial_state_prior=initial_state_prior,
        initial_step=initial_step)


# Copyright 2018 The TensorFlow Probability Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://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.
# ============================================================================
"""The `JointDistributionNamed` class."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections

from tensorflow_probability.python.distributions import joint_distribution_sequential
from tensorflow_probability.python.internal import distribution_util

__all__ = [
    'JointDistributionNamed',
]

class JointDistributionNamed(
    joint_distribution_sequential.JointDistributionSequential):
  """Joint distribution parameterized by named distribution-making functions.

  This distribution enables both sampling and joint probability computation from
  a single model specification.

  A joint distribution is a collection of possibly interdependent distributions.
  Like `JointDistributionSequential`, `JointDistributionNamed` is parameterized
  by several distribution-making functions. Unlike `JointDistributionNamed`,
  each distribution-making function must have its own key. Additionally every
  distribution-making function's arguments must refer to only specified keys.

  #### Mathematical Details

  Internally `JointDistributionNamed` implements the chain rule of probability.
  That is, the probability function of a length-`d` vector `x` is,

  ```none
  p(x) = prod{ p(x[i] | x[:i]) : i = 0, ..., (d - 1) }
  ```

  The `JointDistributionNamed` is parameterized by a `dict` (or `namedtuple`)
  composed of either:

  1. `tfp.distributions.Distribution`-like instances or,
  2. `callable`s which return a `tfp.distributions.Distribution`-like instance.

  The "conditioned on" elements are represented by the `callable`'s required
  arguments; every argument must correspond to a key in the named
  distribution-making functions. Distribution-makers which are directly a
  `Distribution`-like instance are allowed for convenience and semantically
  identical a zero argument `callable`. When the maker takes no arguments it is
  preferable to directly provide the distribution instance.

  #### Examples

  ```python
  tfd = tfp.distributions

  # Consider the following generative model:
  #     e ~ Exponential(rate=[100,120])
  #     g ~ Gamma(concentration=e[0], rate=e[1])
  #     n ~ Normal(loc=0, scale=2.)
  #     m ~ Normal(loc=n, scale=g)
  #     for i = 1, ..., 12:
  #       x[i] ~ Bernoulli(logits=m)

  # In TFP, we can write this as:
  joint = tfd.JointDistributionNamed(dict(
      e=             tfd.Independent(tfd.Exponential(rate=[100, 120]), 1),
      g=lambda    e: tfd.Gamma(concentration=e[..., 0], rate=e[..., 1]),
      n=             tfd.Normal(loc=0, scale=2.),
      m=lambda n, g: tfd.Normal(loc=n, scale=g),
      x=lambda    m: tfd.Sample(tfd.Bernoulli(logits=m), 12),
  ))
  # Notice the 1:1 correspondence between "math" and "code". Further, notice
  # that unlike `JointDistributionSequential`, there is no need to put the
  # distribution-making functions in topologically sorted order nor is it ever
  # necessary to use dummy arguments to skip dependencies.

  x = joint.sample()
  # ==> A 5-element `dict` of tfd.Distribution instances.
  joint.log_prob(x)
  # ==> A scalar `Tensor` representing the total log prob under all five
  #     distributions.

  joint._resolve_graph()
  # ==> (('e', ()),
  #      ('g', ('e',)),
  #      ('n', ()),
  #      ('m', ('n', 'g')),
  #      ('x', ('m',)))
  ```

  #### Discussion

  `JointDistributionNamed` topologically sorts the distribution-making functions
  and calls each by feeding in all previously created dependencies. A
  distribution-maker must either be a:

  1. `tfd.Distribution`-like instance (e.g., `e` and `n` in the above example),
  2. Python `callable` (e.g., `g`, `m`, `x` in the above example).

  Regarding #1, an object is deemed "`tfd.Distribution`-like" if it has a
  `sample`, `log_prob`, and distribution properties, e.g., `batch_shape`,
  `event_shape`, `dtype`.

  Regarding #2, in addition to using a function (or `lambda`), supplying a TFD
  "`class`" is also permissible, this also being a "Python `callable`." For
  example, instead of writing:
  `lambda loc, scale: tfd.Normal(loc=loc, scale=scale)`
  one could have simply written `tfd.Normal`.

  Notice that directly providing a `tfd.Distribution`-like instance means there
  cannot exist a (dynamic) dependency on other distributions; it is
  "independent" both "computationally" and "statistically." The same is
  self-evidently true of zero-argument `callable`s.

  A distribution instance depends on other distribution instances through the
  distribution making function's *required arguments*. The distribution makers'
  arguments are parameterized by samples from the corresponding previously
  constructed distributions. ("Previous" in the sense of a topological sorting
  of dependencies.)

  **Note**: unlike other non-`JointDistribution` distributions in
  `tfp.distributions`, `JointDistribution.sample` (and subclasses) return a
  structure of  `Tensor`s rather than a `Tensor`.  A structure can be a `list`,
  `tuple`, `dict`, `collections.namedtuple`, etc. Accordingly
  `joint.batch_shape` returns a structure of `TensorShape`s for each of the
  distributions' batch shapes and `joint.batch_shape_tensor()` returns a
  structure of `Tensor`s for each of the distributions' event shapes. (Same with
  `event_shape` analogues.)

  **Note**: unlike other non-`JointDistribution` distributions in
  `tfp.distributions`, `JointDistributionNamed.sample` (and subclasses) return a
  structure of  `Tensor`s rather than a `Tensor`.  A structure can be anything
  which is convertible to `dict`. This can be a `dict`,
  `collections.namedtuple`, etc. Accordingly `joint.batch_shape` returns a
  structure of `TensorShape`s for each of the distributions' batch shapes and
  `joint.batch_shape_tensor()` returns a structure of `Tensor`s for each of the
  distributions' event shapes. (Same with `event_shape` analogues.)
  """

  def __init__(self, model, validate_args=False, name=None):
    """Construct the `JointDistributionNamed` distribution.

    Args:
      model: Python `dict` or `namedtuple` of distribution-making functions each
        with required args corresponding only to other keys.
      validate_args: Python `bool`.  Whether to validate input with asserts.
        If `validate_args` is `False`, and the inputs are invalid,
        correct behavior is not guaranteed.
        Default value: `False`.
      name: The name for ops managed by the distribution.
        Default value: `None` (i.e., `"JointDistributionNamed"`).
    """
    super(JointDistributionNamed, self).__init__(
        model, validate_args, name or 'JointDistributionNamed')

  def _build(self, model):
    """Creates `dist_fn`, `dist_fn_wrapped`, `dist_fn_args`, `dist_fn_name`."""
    if not _is_dict_like(model):
      raise TypeError('`model` must be convertible to `dict` (saw: {}).'.format(
          type(model).__name__))
    [
        self._dist_fn,
        self._dist_fn_wrapped,
        self._dist_fn_args,
        self._dist_fn_name,  # JointDistributionSequential doesn't have this.
    ] = _prob_chain_rule_flatten(model)

  def _unflatten(self, xs):
    kwargs = dict(zip(self._dist_fn_name, tuple(xs)))
    return type(self.model)(**kwargs)

  def _flatten(self, xs):
    if xs is None:
      return (None,) * len(self._dist_fn_name)
    if hasattr(xs, 'get'):
      return tuple(xs.get(n, None) for n in self._dist_fn_name)
    return tuple(getattr(xs, n) for n in self._dist_fn_name)

class _Node(object):

  def __init__(self, name, parents):
    self.name = name
    self.parents = parents
    self.depth = -1

def _depth(g):
  """Computes the number of edges on longest path from node to root."""
  def _explore(v):
    if v.depth < 0:
      v.depth = ((1 + max([-1] + [_explore(annotated_graph[u])
                                  for u in v.parents]))
                 if v.parents else 0)
    return v.depth
  annotated_graph = {k: _Node(k, v) for k, v in g.items()}
  for v in annotated_graph.values():
    _explore(v)
  return annotated_graph

def _best_order(g):
  """Creates tuple of str tuple-str pairs representing resolved & sorted DAG."""
  def _explore(u):
    """Recursive function to ascend up through unvisited dependencies."""
    if u.depth < 0:
      return  # Already visited.
    if not u.parents:
      result.append((u.name, u.parents))
      u.depth = -1  # Mark visited.
      return
    b = (u.name, [])
    result.append(b)
    u.depth = -1  # Mark visited.
    d = 0
    for v in sorted((g.get(p) for p in u.parents), key=lambda v: v.depth):
      n0 = len(result)
      _explore(v)
      n1 = len(result)
      b[1].extend(['_']*d + [v.name])
      d = n1 - n0 - 1
  g = _depth(g)
  result = []
  for u in sorted(g.values(), key=lambda v: v.depth, reverse=True):
    _explore(u)
  return tuple(reversed(result))

def _prob_chain_rule_flatten(named_makers):
  """Creates lists of callables suitable for JDSeq."""
  def _make(dist_fn, args):
    if args is None:
      return lambda *_: dist_fn
    if not args:
      return lambda *_: dist_fn()
    def _fn(*xs):
      kwargs = dict(zip(args, reversed(xs[-len(args):])))
      kwargs.pop('_', None)
      return dist_fn(**kwargs)
    return _fn
  named_makers = _convert_to_dict(named_makers)
  g = {k: (None if distribution_util.is_distribution_instance(v)
           else joint_distribution_sequential._get_required_args(v))  # pylint: disable=protected-access
       for k, v in named_makers.items()}
  g = _best_order(g)
  dist_fn_name, dist_fn_args = zip(*g)
  dist_fn_args = tuple(None if a is None else tuple(a) for a in dist_fn_args)
  dist_fn_wrapped = tuple(_make(named_makers[name], parents)
                          for (name, parents) in g)
  dist_fn = tuple(named_makers.get(n) for n in dist_fn_name)
  return dist_fn, dist_fn_wrapped, dist_fn_args, dist_fn_name

def _is_dict_like(x):
  """Returns `True` if input is convertible to `dict`, `False` otherwise."""
  if hasattr(x, '_asdict'):
    return True
  return isinstance(x, collections.Mapping)

def _convert_to_dict(x):
  """Converts input to `dict`."""
  if hasattr(x, '_asdict'):
    return x._asdict()
  return dict(x)