docs: def block_diag(*blocks: np.ndarray) -> np.ndarray:
      """Concatenates blocks into a block diagonal matrix.
  
      Args:
          *blocks: Square matrices to place along the diagonal of the result.
  
      Returns:
          A block diagonal matrix with the given blocks along its diagonal.
  
      Raises:
          ValueError: A block isn't square.
      """
      for b in blocks:
          if b.shape[0] != b.shape[1]:
              raise ValueError('Blocks must be square.')
  
      if not blocks:
          return np.zeros((0, 0), dtype=np.complex128)
  
      n = sum(b.shape[0] for b in blocks)
      dtype = functools.reduce(_merge_dtypes, (b.dtype for b in blocks))
  
      result = np.zeros(shape=(n, n), dtype=dtype)
      i = 0
      for b in blocks:
          j = i + b.shape[0]
          result[i:j, i:j] = b
          i = j
  
      return result