class_docstr: Base TFDS builder.
  
  The functions to overwrite are:
  
  * `_info`: Defines dataset features and metadata (see TFDS docs:
    https://www.tensorflow.org/datasets/add_dataset#_info_dataset_metadata)
  * `split_to_scene_configs`: Define the splits and scenes
  * `generate_scene`: Generate a single scene and returns the examples.
  
  Usage: See `HelloBeamWorker` for an example.
class_docstr_tok: ["Base","TFDS","builder",".","The","functions","to","overwrite","are",":","*","`","_info","`",":","Defines","dataset","features","and","metadata","(","see","TFDS","docs",":","https",":","\/\/www.tensorflow.org\/datasets\/add_dataset","#","_info_dataset_metadata",")","*","`","split_to_scene_configs","`",":","Define","the","splits","and","scenes","*","`","generate_scene","`",":","Generate","a","single","scene","and","returns","the","examples",".","Usage",":","See","`","HelloBeamWorker","`","for","an","example","."]
docstr: Returns the split PCollection.
  
  The pipeline is as follow:
  
  1. Collect all scenes to generate (`split_to_scene_configs`)
  2. Generate each scene in parallel (`_generate_single_scene`)
  3. For each scene, flatten the examples to be distributed across splits (
     `_flatten_split_examples`)
  4. Group examples per split
  
  Args:
    dl_manager: Unused
    pipeline: Root beam pipeline
  
  Returns:
    The dict mapping split names to PCollection of `(key, examples)`
docstr_tok: ["Returns","the","split","PCollection",".","The","pipeline","is","as","follow",":","1",".","Collect","all","scenes","to","generate","(","`","split_to_scene_configs","`",")","2",".","Generate","each","scene","in","parallel","(","`","_generate_single_scene","`",")","3",".","For","each","scene",",","flatten","the","examples","to","be","distributed","across","splits","(","`","_flatten_split_examples","`",")","4",".","Group","examples","per","split","Args",":","dl_manager",":","Unused","pipeline",":","Root","beam","pipeline","Returns",":","The","dict","mapping","split","names","to","PCollection","of","`","(","key",",","examples",")","`"]
code: def _split_generators(
      self,
      dl_manager: tfds.download.DownloadManager,
      pipeline: beam.Pipeline,
  ) -> SplitDict[beam.PCollection[tuple[tfds.typing.Key, ExDict]]]:
      """Returns the split PCollection.
  
      The pipeline is as follow:
  
      1. Collect all scenes to generate (`split_to_scene_configs`)
      2. Generate each scene in parallel (`_generate_single_scene`)
      3. For each scene, flatten the examples to be distributed across splits (
         `_flatten_split_examples`)
      4. Group examples per split
  
      Args:
        dl_manager: Unused
        pipeline: Root beam pipeline
  
      Returns:
        The dict mapping split names to PCollection of `(key, examples)`
      """
      del dl_manager
  
      # Step 1.
      split_to_scenes = SplitScenesMapping(
          self.split_to_scene_configs()
      )
  
      data_per_split = (
          pipeline
          | "Creating scenes"
          >> beam.Create(
              split_to_scenes.scene_id_to_scene_config.items()
          )
          # tuple[<scene_id>, SceneConfig]
          # Step 2.
          | "Generate scenes"
          >> beam.Map(
              self._generate_single_scene,
              split_to_scenes=split_to_scenes,
          )
          # tuple[<key>, SplitToSceneExs]
          # Step 3.
          | "Select split" >> beam.FlatMap(self._flatten_split_examples)
          # tuple[<split_name>, <key>, ExDict]
          # Step 4.
          | "Partition by split"
          >> beam.Partition(
              lambda ex, _: split_to_scenes.split_name_to_split_index[
                  ex[0]
              ],  # pylint: disable=too-many-function-args
              len(split_to_scenes.split_name_to_split_index),
          )
          # Each split is a separate PCollection
          # tuple[<split_name>, <key>, ExDict]
      )
  
      split_to_exs = dict(
          zip(split_to_scenes.split_names, data_per_split)
      )
      # Remove the split name:
      # tuple[<split_name>, <key>, ExDict] -> tuple[<key>, ExDict]
      split_to_exs = {
          split_name: examples
          | f"Remove {split_name}" >> beam.Map(lambda ex: ex[1:])
          for split_name, examples in split_to_exs.items()
      }
  
      # Returns the mapping <split_name> -> <PCollection>
      return split_to_exs
  
imports: from __future__ import annotations
  import abc
  import collections
  import dataclasses
  import itertools
  import tempfile
  from typing import Any, Dict, Iterator, List, Optional, TypeVar, Union
  import apache_beam
  from etils import edc
  from etils import epath
  import kubric
  import tensorflow_datasets
label_desc: Use image node_id 1 for calling the KubricBuilder obj's underlying member method code with example usage: obj._split_generators(dl_manager, pipeline) and returns: split_to_exs