code: def distribute(iterable, layout: ModelLayout):
      """
      Of each group of layout.n_replicas successive items from the iterable, pick the one with index
      `layout.replica_idx`.
      Makes sure that the underlying iterator is advanced at the same pace no matter what replica_idx is.
      """
      it = iter(iterable)
      try:
          while True:
              for i in range(layout.replica_idx):
                  next(it)
              ret = next(it)
              for i in range(
                  layout.n_replicas - layout.replica_idx - 1
              ):
                  next(it)
              yield ret
      except StopIteration:
          return
  
label_desc: Use image node_id 1 for calling a global function with example usage: distribute(iterable, layout) without return types