docstr: Tests Latent Dirichlet Allocation joint model.
  
  The LDA generative process can be written as:
  
  ```none
  N[i] ~ Poisson(xi)
  theta[i] ~ Dirichlet(alpha)
  Z[i] ~ Multinomial(N[i], theta[i])
  for k in 1...K:
    X[i,k] ~ Multinomial(Z[i, k], beta[j])
  ```
  
  Typically `xi` is specified and `alpha`, `beta` are fit using type-II
  maximum likelihood estimators.
  
  Reference: http://www.jmlr.org/papers/volume3/blei03a/blei03a.pdf
docstr_tok: ["Tests","Latent","Dirichlet","Allocation","joint","model",".","The","LDA","generative","process","can","be","written","as",":","``","`","none","N","[","i","]","~","Poisson","(","xi",")","theta","[","i","]","~","Dirichlet","(","alpha",")","Z","[","i","]","~","Multinomial","(","N","[","i","]",",","theta","[","i","]",")","for","k","in","1","...","K",":","X","[","i",",","k","]","~","Multinomial","(","Z","[","i",",","k","]",",","beta","[","j","]",")","``","`","Typically","`","xi","`","is","specified","and","`","alpha","`",",","`","beta","`","are","fit","using","type-II","maximum","likelihood","estimators",".","Reference",":","http",":","\/\/www.jmlr.org\/papers\/volume3\/blei03a\/blei03a.pdf"]
code: def test_latent_dirichlet_allocation(
      self, jd_class
  ):  # pylint: disable=g-doc-args
      """Tests Latent Dirichlet Allocation joint model.
  
      The LDA generative process can be written as:
  
      ```none
      N[i] ~ Poisson(xi)
      theta[i] ~ Dirichlet(alpha)
      Z[i] ~ Multinomial(N[i], theta[i])
      for k in 1...K:
        X[i,k] ~ Multinomial(Z[i, k], beta[j])
      ```
  
      Typically `xi` is specified and `alpha`, `beta` are fit using type-II
      maximum likelihood estimators.
  
      Reference: http://www.jmlr.org/papers/volume3/blei03a/blei03a.pdf
      """
      seed = test_util.test_seed_stream()
      # Hyperparameters.
      num_topics = 3
      num_words = 10
      avg_doc_length = 5
      u = uniform.Uniform(low=-1.0, high=1.0)
      alpha = deferred_tensor.TransformedVariable(
          u.sample([num_topics], seed=seed()),
          softplus.Softplus(),
          name="alpha",
      )
      beta = tf.Variable(
          u.sample([num_topics, num_words], seed=seed()), name="beta"
      )
  
      # Note near 1:1 with mathematical specification. The main distinction is the
      # use of Independent--this lets us easily aggregate multinomials across
      # topics (and in any "shape" of documents).
      def lda_coroutine_model():
          n = yield Root(poisson.Poisson(rate=avg_doc_length))
          theta = yield Root(dirichlet.Dirichlet(concentration=alpha))
          z = yield multinomial.Multinomial(total_count=n, probs=theta)
          yield multinomial.Multinomial(total_count=z, logits=beta)
  
      if jd_class is jdab.JointDistributionCoroutineAutoBatched:
          model = lda_coroutine_model
      elif jd_class is jdab.JointDistributionSequentialAutoBatched:
          model = [
              poisson.Poisson(rate=avg_doc_length),  # n
              dirichlet.Dirichlet(concentration=alpha),  # theta
              # z
              lambda theta, n: multinomial.Multinomial(
                  total_count=n, probs=theta
              ),
              lambda z: multinomial.Multinomial(
                  total_count=z, logits=beta
              ),
          ]
      elif jd_class is jdab.JointDistributionNamedAutoBatched:
          model = collections.OrderedDict(
              (
                  ("n", poisson.Poisson(rate=avg_doc_length)),
                  ("theta", dirichlet.Dirichlet(concentration=alpha)),
                  (
                      "z",
                      lambda theta, n: multinomial.Multinomial(  # pylint: disable=g-long-lambda
                          total_count=n, probs=theta
                      ),
                  ),
                  (
                      "X",
                      lambda z: multinomial.Multinomial(
                          total_count=z, logits=beta
                      ),
                  ),
              )
          )
  
      # TODO(b/159842104): Enable autovectorization for Multinomial sampling.
      lda = jd_class(
          model, validate_args=True, use_vectorized_map=False
      )
  
      # Now, let's sample some "documents" and compute the log-prob of each.
      docs_shape = [2, 4]  # That is, 8 docs in the shape of [2, 4].
      sample = lda.sample(docs_shape, seed=seed())
      log_probs = lda.log_prob(sample)
      self.assertEqual(docs_shape, log_probs.shape)
  
      # Verify we correctly track trainable variables.
      self.assertLen(lda.trainable_variables, 2)
      self.assertIs(
          alpha.pretransformed_input, lda.trainable_variables[0]
      )
      self.assertIs(beta, lda.trainable_variables[1])
  
      # Ensure we can compute gradients.
      with tf.GradientTape() as tape:
          # Note: The samples are not taped, hence implicitly "stop_gradient."
          negloglik = -lda.log_prob(sample)
      grads = tape.gradient(negloglik, lda.trainable_variables)
  
      self.assertLen(grads, 2)
      self.assertAllEqual(
          (alpha.pretransformed_input.shape, beta.shape),
          (grads[0].shape, grads[1].shape),
      )
      self.assertAllNotNone(grads)
  
imports: import collections
  import os
  from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import bijector_test_util
  from tensorflow_probability.python.bijectors import exp
  from tensorflow_probability.python.bijectors import softplus
  from tensorflow_probability.python.distributions import bernoulli
  from tensorflow_probability.python.distributions import deterministic
  from tensorflow_probability.python.distributions import dirichlet
  from tensorflow_probability.python.distributions import exponential
  from tensorflow_probability.python.distributions import half_normal
  from tensorflow_probability.python.distributions import inverse_gamma
  from tensorflow_probability.python.distributions import joint_distribution_auto_batched
  from tensorflow_probability.python.distributions import joint_distribution_coroutine
  from tensorflow_probability.python.distributions import joint_distribution_named
  from tensorflow_probability.python.distributions import joint_distribution_sequential
  from tensorflow_probability.python.distributions import lognormal
  from tensorflow_probability.python.distributions import multinomial
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import poisson
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.distributions import student_t
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.distributions import uniform
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.util import deferred_tensor
label_desc: Use image node_id 18 for calling the JointDistributionAutoBatchedTest obj's underlying member method code with example usage: obj.test_latent_dirichlet_allocation(jd_class) without return types