code: def _to_tfds_split(split: Split) -> tfds.Split:
      """Returns the TFDS split appropriately sharded."""
      if split == Split.TRAIN:
          return tfds.Split.TRAIN
      else:
          assert split == Split.TEST
          return tfds.Split.VALIDATION
  
file_path: objax/examples/image_classification/imagenet_resnet50_data.py
imports: import enum
  from typing import Optional, Sequence, Tuple
  import jax
  import numpy
  import tensorflow.compat.v2
  import tensorflow_datasets
label_desc: Use image node_id 1 for calling a global function with example usage: _to_tfds_split(split) and returns: tfds, tfds