code: def get_num_examples(dataset, split, data_dir=None):
      """Returns the total number of examples in a dataset split."""
      builder = get_builder(dataset, data_dir)
      # Download dataset:
      builder.download_and_prepare()
      num_examples = builder.info.splits[split].num_examples
      remainder = num_examples % jax.process_count()
      if remainder:
          warning = (
              f"Dropping {remainder} examples for the "
              f"{builder.info.name} dataset, {split} split. "
              "The reason is that all hosts should have the same number "
              "of examples in order to guarantee that they stay in sync."
          )
          logging.warning(warning)
  
      return num_examples
  
imports: import collections
  import dataclasses
  import functools
  import itertools
  from typing import Any, Callable, Dict, Iterator, Optional, Sequence, Union
  from absl import logging
  from flax.training import common_utils
  import jax
  import jax.numpy
  import numpy
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 14 for calling a global function with example usage: get_num_examples(dataset, split, data_dir) and returns: num_examples