text: import torch
  from torch_scatter import scatter_add
  
  
  def to_dense_adj(edge_index, batch=None, edge_attr=None, max_num_nodes=None):
      r"""Converts batched sparse adjacency matrices given by edge indices and
      edge attributes to a single dense batched adjacency matrix.
  
      Args:
          edge_index (LongTensor): The edge indices.
          batch (LongTensor, optional): Batch vector
              :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
              node to a specific example. (default: :obj:`None`)
          edge_attr (Tensor, optional): Edge weights or multi-dimensional edge
              features. (default: :obj:`None`)
          max_num_nodes (int, optional): The size of the output node dimension.
              (default: :obj:`None`)
  
      :rtype: :class:`Tensor`
      """
      if batch is None:
          batch = edge_index.new_zeros(edge_index.max().item() + 1)
  
      batch_size = batch[-1].item() + 1
      one = batch.new_ones(batch.size(0))
      num_nodes = scatter_add(one, batch, dim=0, dim_size=batch_size)
      cum_nodes = torch.cat([batch.new_zeros(1), num_nodes.cumsum(dim=0)])
  
      if max_num_nodes is None:
          max_num_nodes = num_nodes.max().item()
  
      size = [batch_size, max_num_nodes, max_num_nodes]
      size = size if edge_attr is None else size + list(edge_attr.size())[1:]
      dtype = torch.float if edge_attr is None else edge_attr.dtype
      adj = torch.zeros(size, dtype=dtype, device=edge_index.device)
  
      edge_index_0 = batch[edge_index[0]].view(1, -1)
      edge_index_1 = edge_index[0] - cum_nodes[batch][edge_index[0]]
      edge_index_2 = edge_index[1] - cum_nodes[batch][edge_index[1]]
  
      if edge_attr is None:
          adj[edge_index_0, edge_index_1, edge_index_2] = 1
      else:
          adj[edge_index_0, edge_index_1, edge_index_2] = edge_attr
  
      return adj
  