# 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.
# ============================================================================
"""Tests for the JointDistributionNamed."""

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

import collections

# Dependency imports

from absl.testing import parameterized
import numpy as np
import tensorflow.compat.v2 as tf
from tensorflow_probability.python import bijectors as tfb
from tensorflow_probability.python import distributions as tfd
from tensorflow_probability.python.internal import test_util


# Defer creating test dists (by hiding them in functions) until we know what
# execution regime (eager/graph/tf-function) the test will run under.
def basic_ordered_model_fn():
  return collections.OrderedDict(
      (('a', tfd.Normal(0., 1.)),
       ('e',
        tfd.Independent(tfd.Exponential(rate=[100, 120]),
                        reinterpreted_batch_ndims=1)),
       ('x', lambda e: tfd.Gamma(concentration=e[..., 0], rate=e[..., 1])),
      ))


def nested_lists_model_fn():
  return collections.OrderedDict((
      ('abc', tfd.JointDistributionSequential([
          tfd.MultivariateNormalDiag([0., 0.], [1., 1.]),
          tfd.JointDistributionSequential(
              [tfd.StudentT(3., -2., 5.),
               tfd.Exponential(4.)])])),
      ('de', lambda abc: tfd.JointDistributionSequential([  # pylint: disable=g-long-lambda
          tfd.Normal(abc[0] * abc[1][0], abc[1][1]),
          tfd.Normal(abc[0] + abc[1][0], abc[1][1])]))))


@test_util.test_all_tf_execution_regimes
class JointDistributionNamedTest(test_util.TestCase):

  def test_dict_sample_log_prob(self):
    # pylint: disable=bad-whitespace
    d = tfd.JointDistributionNamed(dict(
        e    =          tfd.Independent(tfd.Exponential(rate=[100, 120]), 1),
        scale=lambda e: tfd.Gamma(concentration=e[..., 0], rate=e[..., 1]),
        loc  =          tfd.Normal(loc=0, scale=2.),
        m    =          tfd.Normal,
        x    =lambda m: tfd.Sample(tfd.Bernoulli(logits=m), 12)),
                                   validate_args=True)
    # pylint: enable=bad-whitespace

    self.assertEqual(
        (
            ('e', ()),
            ('scale', ('e',)),
            ('loc', ()),
            ('m', ('loc', 'scale')),
            ('x', ('m',)),
        ),
        d.resolve_graph())

    xs = d.sample(seed=test_util.test_seed())
    self.assertLen(xs, 5)
    # We'll verify the shapes work as intended when we plumb these back into the
    # respective log_probs.

    ds, _ = d.sample_distributions(value=xs, seed=test_util.test_seed())
    self.assertLen(ds, 5)
    self.assertIsInstance(ds['e'], tfd.Independent)
    self.assertIsInstance(ds['scale'], tfd.Gamma)
    self.assertIsInstance(ds['loc'], tfd.Normal)
    self.assertIsInstance(ds['m'], tfd.Normal)
    self.assertIsInstance(ds['x'], tfd.Sample)

    # Static properties.
    self.assertAllEqual(
        {'e': tf.float32, 'scale': tf.float32, 'loc': tf.float32,
         'm': tf.float32, 'x': tf.int32},
        d.dtype)

    batch_shape_tensor_, event_shape_tensor_ = self.evaluate([
        d.batch_shape_tensor(), d.event_shape_tensor()])

    expected_batch_shape = {
        'e': [], 'scale': [], 'loc': [], 'm': [], 'x': []}
    batch_tensorshape = d.batch_shape
    for k in expected_batch_shape:
      self.assertAllEqual(expected_batch_shape[k], batch_tensorshape[k])
      self.assertAllEqual(expected_batch_shape[k], batch_shape_tensor_[k])

    expected_event_shape = {
        'e': [2], 'scale': [], 'loc': [], 'm': [], 'x': [12]}
    event_tensorshape = d.event_shape
    for k in expected_event_shape:
      self.assertAllEqual(expected_event_shape[k], event_tensorshape[k])
      self.assertAllEqual(expected_event_shape[k], event_shape_tensor_[k])

    expected_jlp = sum(ds[k].log_prob(xs[k]) for k in ds.keys())
    actual_jlp = d.log_prob(xs)
    self.assertAllClose(*self.evaluate([expected_jlp, actual_jlp]),
                        atol=0., rtol=1e-4)

  def test_namedtuple_sample_log_prob(self):
    Model = collections.namedtuple('Model', ['e', 'scale', 'loc', 'm', 'x'])  # pylint: disable=invalid-name
    # pylint: disable=bad-whitespace
    model = Model(
        e    =          tfd.Independent(tfd.Exponential(rate=[100, 120]), 1),
        scale=lambda e: tfd.Gamma(concentration=e[..., 0], rate=e[..., 1]),
        loc  =          tfd.Normal(loc=0, scale=2.),
        m    =          tfd.Normal,
        x    =lambda m: tfd.Sample(tfd.Bernoulli(logits=m), 12))
    # pylint: enable=bad-whitespace
    d = tfd.JointDistributionNamed(model, validate_args=True)

    self.assertEqual(
        (
            ('e', ()),
            ('scale', ('e',)),
            ('loc', ()),
            ('m', ('loc', 'scale')),
            ('x', ('m',)),
        ),
        d.resolve_graph())

    xs = d.sample(seed=test_util.test_seed())
    self.assertLen(xs, 5)
    # We'll verify the shapes work as intended when we plumb these back into the
    # respective log_probs.

    ds, _ = d.sample_distributions(value=xs, seed=test_util.test_seed())
    self.assertLen(ds, 5)
    self.assertIsInstance(ds.e, tfd.Independent)
    self.assertIsInstance(ds.scale, tfd.Gamma)
    self.assertIsInstance(ds.loc, tfd.Normal)
    self.assertIsInstance(ds.m, tfd.Normal)
    self.assertIsInstance(ds.x, tfd.Sample)

    # Static properties.
    self.assertAllEqual(Model(e=tf.float32, scale=tf.float32, loc=tf.float32,
                              m=tf.float32, x=tf.int32),
                        d.dtype)

    batch_shape_tensor_, event_shape_tensor_ = self.evaluate([
        d.batch_shape_tensor(), d.event_shape_tensor()])

    expected_batch_shape = Model(e=[], scale=[], loc=[], m=[], x=[])
    for (expected, actual_tensorshape, actual_shape_tensor_) in zip(
        expected_batch_shape, d.batch_shape, batch_shape_tensor_):
      self.assertAllEqual(expected, actual_tensorshape)
      self.assertAllEqual(expected, actual_shape_tensor_)

    expected_event_shape = Model(e=[2], scale=[], loc=[], m=[], x=[12])
    for (expected, actual_tensorshape, actual_shape_tensor_) in zip(
        expected_event_shape, d.event_shape, event_shape_tensor_):
      self.assertAllEqual(expected, actual_tensorshape)
      self.assertAllEqual(expected, actual_shape_tensor_)

    expected_jlp = sum(d.log_prob(x) for d, x in zip(ds, xs))
    actual_jlp = d.log_prob(xs)
    self.assertAllClose(*self.evaluate([expected_jlp, actual_jlp]),
                        atol=0., rtol=1e-4)

  def test_ordereddict_sample_log_prob(self):
    build_ordereddict = lambda e, scale, loc, m, x: collections.OrderedDict([  # pylint: disable=g-long-lambda
        ('e', e), ('scale', scale), ('loc', loc), ('m', m), ('x', x)])

    # pylint: disable=bad-whitespace
    model = build_ordereddict(
        e    =          tfd.Independent(tfd.Exponential(rate=[100, 120]), 1),
        scale=lambda e: tfd.Gamma(concentration=e[..., 0], rate=e[..., 1]),
        loc  =          tfd.Normal(loc=0, scale=2.),
        m    =          tfd.Normal,
        x    =lambda m: tfd.Sample(tfd.Bernoulli(logits=m), 12))
    # pylint: enable=bad-whitespace
    d = tfd.JointDistributionNamed(model, validate_args=True)

    self.assertEqual(
        (
            ('e', ()),
            ('scale', ('e',)),
            ('loc', ()),
            ('m', ('loc', 'scale')),
            ('x', ('m',)),
        ),
        d.resolve_graph())

    xs = d.sample(seed=test_util.test_seed())
    self.assertLen(xs, 5)
    # We'll verify the shapes work as intended when we plumb these back into the
    # respective log_probs.

    ds, _ = d.sample_distributions(value=xs, seed=test_util.test_seed())
    self.assertLen(ds, 5)
    values = tuple(ds.values())
    self.assertIsInstance(values[0], tfd.Independent)
    self.assertIsInstance(values[1], tfd.Gamma)
    self.assertIsInstance(values[2], tfd.Normal)
    self.assertIsInstance(values[3], tfd.Normal)
    self.assertIsInstance(values[4], tfd.Sample)

    # Static properties.
    self.assertAllEqual(build_ordereddict(
        e=tf.float32, scale=tf.float32, loc=tf.float32,
        m=tf.float32, x=tf.int32), d.dtype)

    batch_shape_tensor_, event_shape_tensor_ = self.evaluate([
        d.batch_shape_tensor(), d.event_shape_tensor()])

    expected_batch_shape = build_ordereddict(e=[], scale=[], loc=[], m=[], x=[])
    for (expected, actual_tensorshape, actual_shape_tensor_) in zip(
        expected_batch_shape, d.batch_shape, batch_shape_tensor_):
      self.assertAllEqual(expected, actual_tensorshape)
      self.assertAllEqual(expected, actual_shape_tensor_)

    expected_event_shape = build_ordereddict(
        e=[2], scale=[], loc=[], m=[], x=[12])
    for (expected, actual_tensorshape, actual_shape_tensor_) in zip(
        expected_event_shape, d.event_shape, event_shape_tensor_):
      self.assertAllEqual(expected, actual_tensorshape)
      self.assertAllEqual(expected, actual_shape_tensor_)

    expected_jlp = sum(d.log_prob(x) for d, x in zip(ds.values(), xs.values()))
    actual_jlp = d.log_prob(xs)
    self.assertAllClose(*self.evaluate([expected_jlp, actual_jlp]),
                        atol=0., rtol=1e-4)

  def test_can_call_log_prob_with_kwargs(self):

    d = tfd.JointDistributionNamed({
        'e': tfd.Normal(0., 1.),
        'a': tfd.Independent(
            tfd.Exponential(rate=[100, 120]),
            reinterpreted_batch_ndims=1),
        'x': lambda a: tfd.Gamma(concentration=a[..., 0], rate=a[..., 1])
    }, validate_args=True)

    sample = self.evaluate(d.sample([2, 3], seed=test_util.test_seed()))
    e, a, x = sample['e'], sample['a'], sample['x']

    lp_value_positional = self.evaluate(d.log_prob({'e': e, 'a': a, 'x': x}))
    lp_value_named = self.evaluate(d.log_prob(value={'e': e, 'a': a, 'x': x}))
    # Assert all close (rather than equal) because order is not defined for
    # dicts, and reordering the computation can give subtly different results.
    self.assertAllClose(lp_value_positional, lp_value_named)

    lp_kwargs = self.evaluate(d.log_prob(a=a, e=e, x=x))
    self.assertAllClose(lp_value_positional, lp_kwargs)

    with self.assertRaisesRegexp(ValueError,
                                 'Joint distribution with unordered variables '
                                 "can't take positional args"):
      lp_kwargs = d.log_prob(e, a, x)

  @parameterized.named_parameters(
      ('_sample', lambda d, **kwargs: d.sample(**kwargs)),
      ('_sample_and_log_prob',
       lambda d, **kwargs: d.experimental_sample_and_log_prob(**kwargs)[0]),
  )
  def test_nested_partial_value(self, sample_fn):
    innermost = tfd.JointDistributionNamed({
        'a': tfd.Exponential(1.),
        'b': lambda a: tfd.Sample(tfd.LogNormal(a, a), [5]),
    })

    inner = tfd.JointDistributionNamed({
        'c': tfd.Exponential(1.),
        'd': innermost,
    })

    outer = tfd.JointDistributionNamed({
        'e': tfd.Exponential(1.),
        'f': inner,
    })

    seed = test_util.test_seed(sampler_type='stateless')
    true_xs = outer.sample(seed=seed)

    def _update(dict_, **kwargs):
      dict_.copy().update(**kwargs)
      return dict_

    # These asserts work because we advance the stateless seed inside the model
    # whether or not a sample is actually generated.
    partial_xs = _update(true_xs, f=None)
    xs = sample_fn(outer, value=partial_xs, seed=seed)
    self.assertAllCloseNested(true_xs, xs)

    partial_xs = _update(true_xs, e=None)
    xs = sample_fn(outer, value=partial_xs, seed=seed)
    self.assertAllCloseNested(true_xs, xs)

    partial_xs = _update(true_xs, f=_update(true_xs['f'], d=None))
    xs = sample_fn(outer, value=partial_xs, seed=seed)
    self.assertAllCloseNested(true_xs, xs)

    partial_xs = _update(
        true_xs, f=_update(true_xs['f'], d=_update(true_xs['f']['d'], a=None)))
    xs = sample_fn(outer, value=partial_xs, seed=seed)
    self.assertAllCloseNested(true_xs, xs)

  @parameterized.named_parameters(
      ('basic', basic_ordered_model_fn),
      ('nested_lists', nested_lists_model_fn))
  def test_can_call_ordereddict_log_prob_with_args_and_kwargs(self, model_fn):
    # With an OrderedDict, we can pass keyword and/or positional args.
    d = tfd.JointDistributionNamed(model_fn(), validate_args=True)

    # Destructure vector-valued Tensors into Python lists, to mimic the values
    # a user might type.
    def _convert_ndarray_to_list(x):
      if isinstance(x, np.ndarray) and x.ndim > 0:
        return list(x)
      return x
    sample = tf.nest.map_structure(
        _convert_ndarray_to_list,
        self.evaluate(d.sample(seed=test_util.test_seed())))
    sample_dict = dict(sample)

    lp_value_positional = self.evaluate(d.log_prob(sample_dict))
    lp_value_named = self.evaluate(d.log_prob(value=sample_dict))
    self.assertAllClose(lp_value_positional, lp_value_named)

    lp_args = self.evaluate(d.log_prob(*sample.values()))
    self.assertAllClose(lp_value_positional, lp_args)

    lp_kwargs = self.evaluate(d.log_prob(**sample_dict))
    self.assertAllClose(lp_value_positional, lp_kwargs)

    lp_args_then_kwargs = self.evaluate(d.log_prob(
        *list(sample.values())[:1], **dict(list(sample.items())[1:])))
    self.assertAllClose(lp_value_positional, lp_args_then_kwargs)

  def test_can_call_namedtuple_log_prob_with_args_and_kwargs(self):
    # With an namedtuple, we can pass keyword and/or positional args.
    Model = collections.namedtuple('Model', ['e', 'a', 'x'])  # pylint: disable=invalid-name
    d = tfd.JointDistributionNamed(
        Model(e=tfd.Normal(0., 1.),
              a=tfd.Independent(
                  tfd.Exponential(rate=[100, 120]),
                  reinterpreted_batch_ndims=1),
              x=lambda a: tfd.Gamma(concentration=a[..., 0], rate=a[..., 1])),
        validate_args=True)

    sample = self.evaluate(d.sample([2, 3], seed=test_util.test_seed()))
    e, a, x = sample.e, sample.a, sample.x

    lp_value_positional = self.evaluate(d.log_prob(Model(e=e, a=a, x=x)))
    lp_value_named = self.evaluate(d.log_prob(value=Model(e=e, a=a, x=x)))
    self.assertAllClose(lp_value_positional, lp_value_named)

    lp_kwargs = self.evaluate(d.log_prob(e=e, a=a, x=x))
    self.assertAllClose(lp_value_positional, lp_kwargs)

    lp_args = self.evaluate(d.log_prob(e, a, x))
    self.assertAllClose(lp_value_positional, lp_args)

    lp_args_then_kwargs = self.evaluate(d.log_prob(e, a=a, x=x))
    self.assertAllClose(lp_value_positional, lp_args_then_kwargs)

  def test_kl_divergence(self):
    d0 = tfd.JointDistributionNamed(
        dict(e=tfd.Independent(tfd.Exponential(rate=[100, 120]), 1),
             x=tfd.Normal(loc=0, scale=2.)),
        validate_args=True)
    d1 = tfd.JointDistributionNamed(
        dict(e=tfd.Independent(tfd.Exponential(rate=[10, 12]), 1),
             x=tfd.Normal(loc=1, scale=1.)),
        validate_args=True)
    self.assertEqual(d0.model.keys(), d1.model.keys())
    expected_kl = sum(tfd.kl_divergence(d0.model[k], d1.model[k])
                      for k in d0.model.keys())
    actual_kl = tfd.kl_divergence(d0, d1)
    other_actual_kl = d0.kl_divergence(d1)
    expected_kl_, actual_kl_, other_actual_kl_ = self.evaluate([
        expected_kl, actual_kl, other_actual_kl])
    self.assertNear(expected_kl_, actual_kl_, err=1e-5)
    self.assertNear(expected_kl_, other_actual_kl_, err=1e-5)

  def test_cross_entropy(self):
    d0 = tfd.JointDistributionNamed(
        dict(e=tfd.Independent(tfd.Exponential(rate=[100, 120]), 1),
             x=tfd.Normal(loc=0, scale=2.)),
        validate_args=True)
    d1 = tfd.JointDistributionNamed(
        dict(e=tfd.Independent(tfd.Exponential(rate=[10, 12]), 1),
             x=tfd.Normal(loc=1, scale=1.)),
        validate_args=True)
    self.assertEqual(d0.model.keys(), d1.model.keys())
    expected_xent = sum(d0.model[k].cross_entropy(d1.model[k])
                        for k in d0.model.keys())
    actual_xent = d0.cross_entropy(d1)
    expected_xent_, actual_xent_ = self.evaluate([expected_xent, actual_xent])
    self.assertNear(actual_xent_, expected_xent_, err=1e-5)

  def test_norequired_args_maker(self):
    """Test that only non-default args are passed through."""
    with self.assertRaisesWithPredicateMatch(
        ValueError, 'Must pass probs or logits, but not both.'):
      tfd.JointDistributionNamed(dict(logits=tfd.Normal(0., 1.),
                                      x=tfd.Bernoulli))

  def test_dist_fn_takes_kwargs(self):
    dist = tfd.JointDistributionNamed(
        {'positive': tfd.Exponential(rate=1.),
         'negative': tfb.Scale(-1.)(tfd.Exponential(rate=1.)),
         'b': lambda **kwargs: tfd.Normal(loc=kwargs['negative'],  # pylint: disable=g-long-lambda
                                          scale=kwargs['positive'],
                                          validate_args=True),
         'a': lambda **kwargs: tfb.Scale(kwargs['b'])(  # pylint: disable=g-long-lambda
             tfd.Gamma(concentration=-kwargs['negative'],
                       rate=kwargs['positive'],
                       validate_args=True))
         }, validate_args=True)
    lp = dist.log_prob(dist.sample(5, seed=test_util.test_seed()))
    self.assertAllEqual(lp.shape, [5])

  def test_graph_resolution(self):
    # pylint: disable=bad-whitespace
    d = tfd.JointDistributionNamed(dict(
        e    =          tfd.Independent(tfd.Exponential(rate=[100, 120]), 1),
        scale=lambda e: tfd.Gamma(concentration=e[..., 0], rate=e[..., 1]),
        s    =          tfd.HalfNormal(2.5),
        loc  =lambda s: tfd.Normal(loc=0, scale=s),
        df   =          tfd.Exponential(2),
        x    =          tfd.StudentT),
                                   validate_args=True)
    # pylint: enable=bad-whitespace
    self.assertEqual(
        (
            ('e', ()),
            ('scale', ('e',)),
            ('s', ()),
            ('loc', ('s',)),
            ('df', ()),
            ('x', ('df', 'loc', 'scale'))
        ),
        d.resolve_graph())

  @parameterized.parameters('mean', 'mode', 'stddev', 'variance')
  def test_summary_statistic(self, attr):
    d = tfd.JointDistributionNamed(dict(logits=tfd.Normal(0., 1.),
                                        x=tfd.Bernoulli(logits=0.)),
                                   validate_args=True)
    expected = {k: getattr(d.model[k], attr)() for k in d.model.keys()}
    actual = getattr(d, attr)()
    self.assertAllEqual(*self.evaluate([expected, actual]))

  @parameterized.parameters(('covariance',))
  def test_notimplemented_summary_statistic(self, attr):
    d = tfd.JointDistributionNamed(dict(logits=tfd.Normal(0., 1.),
                                        x=tfd.Bernoulli(probs=0.5)),
                                   validate_args=True)
    with self.assertRaisesWithPredicateMatch(
        NotImplementedError,
        attr + ' is not implemented: JointDistributionNamed'):
      getattr(d, attr)()

  @parameterized.parameters(
      'log_cdf', 'cdf', 'log_survival_function', 'survival_function')
  def test_notimplemented_evaluative_statistic(self, attr):
    d = tfd.JointDistributionNamed(dict(logits=tfd.Normal(0., 1.),
                                        x=tfd.Bernoulli(probs=0.5)),
                                   validate_args=True)
    with self.assertRaisesWithPredicateMatch(
        NotImplementedError,
        attr + ' is not implemented: JointDistributionNamed'):
      getattr(d, attr)(dict(logits=0., x=0.5))

  def test_notimplemented_quantile(self):
    d = tfd.JointDistributionNamed(dict(logits=tfd.Normal(0., 1.),
                                        x=tfd.Bernoulli(probs=0.5)),
                                   validate_args=True)
    with self.assertRaisesWithPredicateMatch(
        NotImplementedError,
        'quantile is not implemented: JointDistributionNamed'):
      d.quantile(0.5)

  def test_copy(self):
    pgm = dict(logits=tfd.Normal(0., 1.), probs=tfd.Bernoulli(logits=0.5))
    d = tfd.JointDistributionNamed(pgm, validate_args=True)
    d_copy = d.copy()
    self.assertAllEqual(
        {'model': pgm,
         'validate_args': True,
         'name': 'JointDistributionNamed'},
        d_copy.parameters)

  def test_batch_slicing(self):
    # pylint: disable=bad-whitespace
    d = tfd.JointDistributionNamed(
        dict(s=          tfd.Exponential(rate=[10, 12, 14]),
             n=lambda s: tfd.Normal(loc=0, scale=s),
             x=lambda:   tfd.Beta(concentration0=[3, 2, 1], concentration1=1)),
        validate_args=True)
    # pylint: enable=bad-whitespace

    d0, d1 = d[:1], d[1:]
    x0 = d0.sample(seed=test_util.test_seed())
    x1 = d1.sample(seed=test_util.test_seed())

    self.assertLen(x0, 3)
    self.assertEqual([1], x0['s'].shape)
    self.assertEqual([1], x0['n'].shape)
    self.assertEqual([1], x0['x'].shape)

    self.assertLen(x1, 3)
    self.assertEqual([2], x1['s'].shape)
    self.assertEqual([2], x1['n'].shape)
    self.assertEqual([2], x1['x'].shape)

  def test_sample_shape_propagation_default_behavior(self):
    # pylint: disable=bad-whitespace
    d = tfd.JointDistributionNamed(dict(
        e    =          tfd.Independent(tfd.Exponential(rate=[100, 120]), 1),
        scale=lambda e: tfd.Gamma(concentration=e[..., 0], rate=e[..., 1]),
        s    =          tfd.HalfNormal(2.5),
        loc  =lambda s: tfd.Normal(loc=0, scale=s),
        df   =          tfd.Exponential(2),
        x    =          tfd.StudentT),
                                   validate_args=False)
    # pylint: enable=bad-whitespace
    x = d.sample([2, 3], seed=test_util.test_seed())
    self.assertLen(x, 6)
    self.assertEqual((2, 3, 2), x['e'].shape)
    self.assertEqual((2, 3), x['scale'].shape)
    self.assertEqual((2, 3), x['s'].shape)
    self.assertEqual((2, 3), x['loc'].shape)
    self.assertEqual((2, 3), x['df'].shape)
    self.assertEqual((2, 3), x['x'].shape)
    lp = d.log_prob(x)
    self.assertEqual((2, 3), lp.shape)

  def test_sample_complex_dependency(self):
    # pylint: disable=bad-whitespace
    d = tfd.JointDistributionNamed(
        dict(
            y    =          tfd.StudentT,
            x    =          tfd.StudentT,
            df   =          tfd.Exponential(2),
            loc  =lambda s: tfd.Normal(loc=0, scale=s),
            s    =          tfd.HalfNormal(2.5),
            scale=lambda e: tfd.Gamma(concentration=e[..., 0], rate=e[..., 1]),
            e    =          tfd.Independent(tfd.Exponential(rate=[100, 120]), 1)
        ),
        validate_args=False)

    # pylint: enable=bad-whitespace

    self.assertEqual(
        (
            ('e', ()),
            ('scale', ('e',)),
            ('s', ()),
            ('loc', ('s',)),
            ('df', ()),
            ('y', ('df', 'loc', 'scale')),
            ('x', ('df', 'loc', 'scale')),
        ),
        d.resolve_graph())

    x = d.sample(seed=test_util.test_seed())
    self.assertLen(x, 7)

    ds, s = d.sample_distributions(seed=test_util.test_seed())
    self.assertEqual(ds['x'].parameters['df'], s['df'])
    self.assertEqual(ds['x'].parameters['loc'], s['loc'])
    self.assertEqual(ds['x'].parameters['scale'], s['scale'])
    self.assertEqual(ds['y'].parameters['df'], s['df'])
    self.assertEqual(ds['y'].parameters['loc'], s['loc'])
    self.assertEqual(ds['y'].parameters['scale'], s['scale'])

  def test_default_event_space_bijector(self):
    # pylint: disable=bad-whitespace
    d = tfd.JointDistributionNamed(dict(
        e    =          tfd.Independent(tfd.Exponential(rate=[100, 120]), 1),
        scale=lambda e: tfd.Gamma(concentration=e[..., 0], rate=e[..., 1]),
        s    =          tfd.HalfNormal(2.5),
        loc  =lambda s: tfd.Normal(loc=0, scale=s),
        df   =          tfd.Exponential(2),
        x    =          tfd.StudentT),
                                   validate_args=True)
    # pylint: enable=bad-whitespace

    # The event space bijector is inherited from `JointDistributionSequential`
    # and is tested more thoroughly in the tests for that class.
    b = d.experimental_default_event_space_bijector()
    y = self.evaluate(d.sample(seed=test_util.test_seed()))
    y_ = self.evaluate(b.forward(b.inverse(y)))
    self.assertAllClose(y, y_)

    # Verify that event shapes are passed through and flattened/unflattened
    # correctly.
    forward_event_shapes = b.forward_event_shape(d.event_shape)
    inverse_event_shapes = b.inverse_event_shape(d.event_shape)
    self.assertEqual(forward_event_shapes, d.event_shape)
    self.assertEqual(inverse_event_shapes, d.event_shape)

    # Verify that the outputs of other methods have the correct dict structure.
    forward_event_shape_tensors = b.forward_event_shape_tensor(
        d.event_shape_tensor())
    inverse_event_shape_tensors = b.inverse_event_shape_tensor(
        d.event_shape_tensor())
    for item in [forward_event_shape_tensors, inverse_event_shape_tensors]:
      self.assertSetEqual(set(self.evaluate(item).keys()), set(d.model.keys()))

  def test_sample_kwargs(self):
    joint = tfd.JointDistributionNamed(
        dict(
            a=tfd.Normal(0., 1.),
            b=lambda a: tfd.Normal(a, 1.),
            c=lambda a, b: tfd.Normal(a + b, 1.)))

    seed = test_util.test_seed()
    tf.random.set_seed(seed)
    samples = joint.sample(seed=seed, a=1.)
    # Check the first value is actually 1.
    self.assertEqual(1., self.evaluate(samples['a']))

    # Check the sample is reproducible using the `value` argument.
    tf.random.set_seed(seed)
    samples_named = joint.sample(seed=seed, value={'a': 1.})
    self.assertAllEqual(self.evaluate(samples), self.evaluate(samples_named))

    # Make sure to throw an exception if strange keywords are passed.
    expected_error = (
        'Found unexpected keyword arguments. Distribution names are\n'
        'a, b, c\n'
        'but received\n'
        'z\n'
        'These names were invalid:\n'
        'z')
    with self.assertRaisesRegex(ValueError, expected_error):
      joint.sample(seed=seed, z=2.)

    # Raise if value and keywords are passed.
    with self.assertRaisesRegex(
        ValueError, r'Supplied both `value` and keyword arguments .*'):
      joint.sample(seed=seed, a=1., value={'a': 1})


if __name__ == '__main__':
  test_util.main()
