code: def test_gaussian_mrf_log_prob_block_diag(
      sample_shape, batch_shape, num_steps, hidden_dim, obs_dim
  ):
      # Construct a block-diagonal obs dist, so observations are independent of hidden state.
      obs_dist = random_mvn(
          batch_shape + (num_steps,), hidden_dim + obs_dim
      )
      precision = obs_dist.precision_matrix
      precision[..., :hidden_dim, hidden_dim:] = 0
      precision[..., hidden_dim:, :hidden_dim] = 0
      obs_dist = dist.MultivariateNormal(
          obs_dist.loc, precision_matrix=precision
      )
      marginal_obs_dist = dist.MultivariateNormal(
          obs_dist.loc[..., hidden_dim:],
          precision_matrix=precision[..., hidden_dim:, hidden_dim:],
      )
  
      init_dist = random_mvn(batch_shape, hidden_dim)
      trans_dist = random_mvn(
          batch_shape + (num_steps,), hidden_dim + hidden_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)
      expected_log_prob = marginal_obs_dist.log_prob(data).sum(-1)
      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 16 for calling a global function with example usage: test_gaussian_mrf_log_prob_block_diag(sample_shape, batch_shape, num_steps, hidden_dim, obs_dim) without return types