code: def test_gaussian_mrf_log_prob(
      sample_shape, batch_shape, num_steps, hidden_dim, obs_dim
  ):
      init_dist = random_mvn(batch_shape, hidden_dim)
      trans_dist = random_mvn(
          batch_shape + (num_steps,), hidden_dim + hidden_dim
      )
      obs_dist = random_mvn(
          batch_shape + (num_steps,), hidden_dim + obs_dim
      )
      d = dist.GaussianMRF(init_dist, trans_dist, obs_dist)
      data = obs_dist.sample(sample_shape)[..., hidden_dim:]
      assert data.shape == sample_shape + d.shape()
      actual_log_prob = d.log_prob(data)
  
      # Compare against hand-computed density.
      # We will construct enormous unrolled joint gaussians with shapes:
      #       t | 0 1 2 3 1 2 3      T = 3 in this example
      #   ------+-----------------------------------------
      #    init | H
      #   trans | H H H H            H = hidden
      #     obs |   H H H O O O      O = observed
      # and then combine these using gaussian_tensordot().
      T = num_steps
      init = mvn_to_gaussian(init_dist)
      trans = mvn_to_gaussian(trans_dist)
      obs = mvn_to_gaussian(obs_dist)
  
      unrolled_trans = reduce(
          operator.add,
          [
              trans[..., t].event_pad(
                  left=t * hidden_dim, right=(T - t - 1) * hidden_dim
              )
              for t in range(T)
          ],
      )
      unrolled_obs = reduce(
          operator.add,
          [
              obs[..., t].event_pad(
                  left=t * obs.dim(), right=(T - t - 1) * obs.dim()
              )
              for t in range(T)
          ],
      )
      # Permute obs from HOHOHO to HHHOOO.
      perm = torch.cat(
          [torch.arange(hidden_dim) + t * obs.dim() for t in range(T)]
          + [
              torch.arange(obs_dim) + hidden_dim + t * obs.dim()
              for t in range(T)
          ]
      )
      unrolled_obs = unrolled_obs.event_permute(perm)
      unrolled_data = data.reshape(data.shape[:-2] + (T * obs_dim,))
  
      assert init.dim() == hidden_dim
      assert unrolled_trans.dim() == (1 + T) * hidden_dim
      assert unrolled_obs.dim() == T * (hidden_dim + obs_dim)
      logp_h = gaussian_tensordot(init, unrolled_trans, hidden_dim)
      logp_oh = gaussian_tensordot(logp_h, unrolled_obs, T * hidden_dim)
      logp_h += unrolled_obs.marginalize(right=T * obs_dim)
      expected_log_prob = (
          logp_oh.log_density(unrolled_data) - logp_h.event_logsumexp()
      )
      assert_close(actual_log_prob, expected_log_prob)
  
imports: import logging
  import operator
  from functools import reduce
  import opt_einsum
  import pytest
  import torch
  import pyro
  import pyro.distributions
  from pyro.distributions.hmm import _sequential_gamma_gaussian_tensordot, _sequential_logmatmulexp
  from pyro.distributions.util import broadcast_shape
  from pyro.infer import TraceEnum_ELBO, config_enumerate
  from pyro.ops.gamma_gaussian import gamma_and_mvn_to_gamma_gaussian, gamma_gaussian_tensordot, matrix_and_mvn_to_gamma_gaussian
  from pyro.ops.gaussian import gaussian_tensordot, matrix_and_mvn_to_gaussian, mvn_to_gaussian
  from pyro.ops.indexing import Vindex
  from tests.common import assert_close
  from tests.ops.gamma_gaussian import assert_close_gamma_gaussian, random_gamma, random_gamma_gaussian
  from tests.ops.gaussian import random_mvn
label_desc: Use image node_id 15 for calling a global function with example usage: test_gaussian_mrf_log_prob(sample_shape, batch_shape, num_steps, hidden_dim, obs_dim) without return types