code: def test_dirichlet_multinomial(sample_shape, batch_shape):
      concentration = torch.randn(batch_shape + (3,)).exp()
      total = 10
      probs = torch.tensor([0.2, 0.3, 0.5])
      obs = dist.Multinomial(total, probs).sample(
          sample_shape + batch_shape
      )
  
      f = dist.Dirichlet(concentration)
      g = dist.Dirichlet(1 + obs)
      fg, log_normalizer = f.conjugate_update(g)
  
      x = fg.sample(sample_shape)
      assert_close(
          f.log_prob(x) + g.log_prob(x), fg.log_prob(x) + log_normalizer
      )
  
imports: import pytest
  import torch
  import pyro.distributions
  from tests.common import assert_close
label_desc: Use image node_id 2 for calling a global function with example usage: test_dirichlet_multinomial(sample_shape, batch_shape) without return types