code: def stochastic_blockmodel_graph(
      block_sizes: Union[List[int], torch.Tensor],
      edge_probs: Union[List[List[float]], torch.Tensor],
      directed: bool = False,
  ) -> torch.Tensor:
      r"""Returns the :obj:`edge_index` of a stochastic blockmodel graph.
  
      Args:
          block_sizes ([int] or LongTensor): The sizes of blocks.
          edge_probs ([[float]] or FloatTensor): The density of edges going
              from each block to each other block. Must be symmetric if the
              graph is undirected.
          directed (bool, optional): If set to :obj:`True`, will return a
              directed graph. (default: :obj:`False`)
  
      Examples:
          >>> block_sizes = [2, 2, 4]
          >>> edge_probs = [[0.25, 0.05, 0.02],
          ...               [0.05, 0.35, 0.07],
          ...               [0.02, 0.07, 0.40]]
          >>> stochastic_blockmodel_graph(block_sizes, edge_probs,
          ...                             directed=False)
          tensor([[2, 4, 4, 5, 5, 6, 7, 7],
                  [5, 6, 7, 2, 7, 4, 4, 5]])
  
          >>> stochastic_blockmodel_graph(block_sizes, edge_probs,
          ...                             directed=True)
          tensor([[0, 2, 3, 4, 4, 5, 5],
                  [3, 4, 1, 5, 6, 6, 7]])
      """
      size, prob = block_sizes, edge_probs
  
      if not isinstance(size, torch.Tensor):
          size = torch.tensor(size, dtype=torch.long)
      if not isinstance(prob, torch.Tensor):
          prob = torch.tensor(prob, dtype=torch.float)
  
      assert size.dim() == 1
      assert prob.dim() == 2 and prob.size(0) == prob.size(1)
      assert size.size(0) == prob.size(0)
      if not directed:
          assert torch.allclose(prob, prob.t())
  
      node_idx = torch.cat(
          [size.new_full((b,), i) for i, b in enumerate(size)]
      )
      num_nodes = node_idx.size(0)
  
      if directed:
          idx = torch.arange((num_nodes - 1) * num_nodes)
          idx = idx.view(num_nodes - 1, num_nodes)
          idx = idx + torch.arange(1, num_nodes).view(-1, 1)
          idx = idx.view(-1)
          row = idx.div(num_nodes, rounding_mode="floor")
          col = idx % num_nodes
      else:
          row, col = torch.combinations(
              torch.arange(num_nodes), r=2
          ).t()
  
      mask = torch.bernoulli(prob[node_idx[row], node_idx[col]]).to(
          torch.bool
      )
      edge_index = torch.stack([row[mask], col[mask]], dim=0)
  
      if not directed:
          edge_index = to_undirected(edge_index, num_nodes=num_nodes)
  
      return edge_index
  
imports: import warnings
  from typing import List, Union
  import numpy
  import torch
  from torch_geometric.utils import remove_self_loops, to_undirected
label_desc: Use image node_id 2 for calling a global function with example usage: stochastic_blockmodel_graph(block_sizes, edge_probs, directed) and returns: edge_index