code: def test_block_diag_embed(batch_size, block_size):
      m = (
          torch.randn(block_size)
          .unsqueeze(0)
          .expand((batch_size,) + block_size)
      )
      b = block_diag_embed(m)
  
      assert b.shape == (
          batch_size * block_size[0],
          batch_size * block_size[1],
      )
  
      assert_equal(b.sum(), m.sum())
  
      for k in range(batch_size):
          bottom, top = k * block_size[0], (k + 1) * block_size[0]
          left, right = k * block_size[1], (k + 1) * block_size[1]
          assert_equal(b[bottom:top, left:right], m[k])
  
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 1 for calling a global function with example usage: test_block_diag_embed(batch_size, block_size) without return types