code: def batch_edges_group_by_relation_type(
      edges: EdgeList, *, batch_size: int
  ) -> Iterable[EdgeList]:
      """Split the edges in batches that each contain a single relation type
  
      Batches are all of the same size, except possibly the last one for each
      relation type.
      """
      edge_groups = group_by_relation_type(edges)
      num_edges_left_per_group = torch.tensor(
          [len(edges) for edges in edge_groups], dtype=torch.long
      )
  
      while num_edges_left_per_group.sum() > 0:
          idx = int(
              torch.multinomial(num_edges_left_per_group.float(), 1)
          )
          edge_group = edge_groups[idx]
          offset = len(edge_group) - int(num_edges_left_per_group[idx])
          sub_edges = edge_group[offset : offset + batch_size]
          yield sub_edges
          num_edges_left_per_group[idx] -= len(sub_edges)
  
file_path: PyTorch-BigGraph/torchbiggraph/batching.py
imports: import time
  from abc import ABC, abstractmethod
  from typing import Any, Callable, Dict, Iterable, List, Optional
  import torch
  from torchbiggraph.edgelist import EdgeList
  from torchbiggraph.losses import AbstractLossFunction
  from torchbiggraph.model import MultiRelationEmbedder, override_model, Scores
  from torchbiggraph.stats import Stats
  from torchbiggraph.types import LongTensorType
label_desc: Use image node_id 3 for calling a global function with example usage: batch_edges_group_by_relation_type(edges) without return types