code: def test_block_diag(batch_shape, mat_size, block_size):
      mat = torch.randn(batch_shape + (block_size,) + mat_size)
      mat_embed = block_diag_embed(mat)
      mat_embed_diag = block_diagonal(mat_embed, block_size)
      assert_equal(mat_embed_diag, mat)
  
imports: import math
  import numpy
  import pytest
  import scipy.fftpack
  import torch
  import pyro
  from pyro.ops.tensor_utils import block_diag_embed, block_diagonal, convolve, dct, idct, next_fast_len, periodic_cumsum, periodic_features, periodic_repeat, precision_to_scale_tril, repeated_matmul
  from tests.common import assert_close, assert_equal
label_desc: Use image node_id 2 for calling a global function with example usage: test_block_diag(batch_shape, mat_size, block_size) without return types