def _testMVN(self,
               base_distribution_class,
               base_distribution_kwargs,
               batch_shape=(),
               event_shape=(),
               not_implemented_message=None):
    with self.test_session() as sess:
      # Overriding shapes must be compatible w/bijector; most bijectors are
      # batch_shape agnostic and only care about event_ndims.
      # In the case of `Affine`, if we got it wrong then it would fire an
      # exception due to incompatible dimensions.
      batch_shape_pl = array_ops.placeholder(
          dtypes.int32, name="dynamic_batch_shape")
      event_shape_pl = array_ops.placeholder(
          dtypes.int32, name="dynamic_event_shape")
      feed_dict = {batch_shape_pl: np.array(batch_shape, dtype=np.int32),
                   event_shape_pl: np.array(event_shape, dtype=np.int32)}
      fake_mvn_dynamic = self._cls()(
          distribution=base_distribution_class(validate_args=True,
                                               **base_distribution_kwargs),
          bijector=bs.Affine(shift=self._shift, scale_tril=self._tril),
          batch_shape=batch_shape_pl,
          event_shape=event_shape_pl,
          validate_args=True)

      fake_mvn_static = self._cls()(
          distribution=base_distribution_class(validate_args=True,
                                               **base_distribution_kwargs),
          bijector=bs.Affine(shift=self._shift, scale_tril=self._tril),
          batch_shape=batch_shape,
          event_shape=event_shape,
          validate_args=True)

      actual_mean = np.tile(self._shift, [2, 1])  # Affine elided this tile.
      actual_cov = np.matmul(self._tril, np.transpose(self._tril, [0, 2, 1]))

      def actual_mvn_log_prob(x):
        return np.concatenate([
            [stats.multivariate_normal(
                actual_mean[i], actual_cov[i]).logpdf(x[:, i, :])]
            for i in range(len(actual_cov))]).T

      actual_mvn_entropy = np.concatenate([
          [stats.multivariate_normal(
              actual_mean[i], actual_cov[i]).entropy()]
          for i in range(len(actual_cov))])

      self.assertAllEqual([3], fake_mvn_static.event_shape)
      self.assertAllEqual([2], fake_mvn_static.batch_shape)

      self.assertAllEqual(tensor_shape.TensorShape(None),
                          fake_mvn_dynamic.event_shape)
      self.assertAllEqual(tensor_shape.TensorShape(None),
                          fake_mvn_dynamic.batch_shape)

      x = fake_mvn_static.sample(5, seed=0).eval()
      for unsupported_fn in (fake_mvn_static.log_cdf,
                             fake_mvn_static.cdf,
                             fake_mvn_static.survival_function,
                             fake_mvn_static.log_survival_function):
        with self.assertRaisesRegexp(NotImplementedError,
                                     not_implemented_message):
          unsupported_fn(x)

      num_samples = 5e3
      for fake_mvn, feed_dict in ((fake_mvn_static, {}),
                                  (fake_mvn_dynamic, feed_dict)):
        # Ensure sample works by checking first, second moments.
        y = fake_mvn.sample(int(num_samples), seed=0)
        x = y[0:5, ...]
        sample_mean = math_ops.reduce_mean(y, 0)
        centered_y = array_ops.transpose(y - sample_mean, [1, 2, 0])
        sample_cov = math_ops.matmul(
            centered_y, centered_y, transpose_b=True) / num_samples
        [
            sample_mean_,
            sample_cov_,
            x_,
            fake_event_shape_,
            fake_batch_shape_,
            fake_log_prob_,
            fake_prob_,
            fake_entropy_,
        ] = sess.run([
            sample_mean,
            sample_cov,
            x,
            fake_mvn.event_shape_tensor(),
            fake_mvn.batch_shape_tensor(),
            fake_mvn.log_prob(x),
            fake_mvn.prob(x),
            fake_mvn.entropy(),
        ], feed_dict=feed_dict)

        self.assertAllClose(actual_mean, sample_mean_, atol=0.1, rtol=0.1)
        self.assertAllClose(actual_cov, sample_cov_, atol=0., rtol=0.1)

        # Ensure all other functions work as intended.
        self.assertAllEqual([5, 2, 3], x_.shape)
        self.assertAllEqual([3], fake_event_shape_)
        self.assertAllEqual([2], fake_batch_shape_)
        self.assertAllClose(actual_mvn_log_prob(x_), fake_log_prob_,
                            atol=0., rtol=1e-6)
        self.assertAllClose(np.exp(actual_mvn_log_prob(x_)), fake_prob_,
                            atol=0., rtol=1e-5)
        self.assertAllClose(actual_mvn_entropy, fake_entropy_,
                            atol=0., rtol=1e-6)