code: def get_dataset(
      *,
      batch_size,
      eval_batch_size,
      num_shards,
      dtype_str="float32",
      shuffle_seed=0,
      rng=None,
      prefetch_buffer_size=2,
      dataset_configs=None,
      dataset_service_address: Optional[str] = None,
      seq_length=2048,
  ):
      """Returns generators for the ShapeNet train, validation, and test sets.
  
      Args:
        batch_size: int; Determines the train batch size.
        eval_batch_size: int; Determines the evaluation batch size.
        num_shards: int;  Number of shards --> batch shape: [num_shards, bs, ...].
        dtype_str: Data type of the image (e.g. 'float32').
        shuffle_seed: int; Seed for shuffling the training data.
        rng: JAX rng key, which can be used for augmentation, shuffling, etc.
        prefetch_buffer_size: int; Buffer size for the device prefetch.
        dataset_configs: dict; Dataset specific configurations.
        dataset_service_address: If set, will distribute the training dataset using
          the given tf.data service at the given address.
        seq_length: maximum sequence length.
  
      Returns:
        A dataset_utils.Dataset() which includes a train_iter, a valid_iter,
          a test_iter, and a dict of meta_data.
      """
      shapenet_config = get_dataset_config()
      dataset_configs = dataset_configs or {}
      del rng
      data_augmentations = dataset_configs.get(
          "data_augmentations", ["default"]
      )
      # TODO(dehghani): add mixup data augmentation.
      for da in data_augmentations:
          if da not in ["default"]:
              raise ValueError(
                  f"Data augmentation {data_augmentations} is not "
                  f"(yet) supported in the ShapeNet dataset."
              )
      dtype = getattr(tf, dtype_str)
      onehot_labels = (
          False  # dataset_configs.get('onehot_labels', False)
      )
  
      logging.info("Loading train split of the ShapeNet dataset.")
      train_ds, _ = shapenet_load_split(
          batch_size,
          shapenet_config,
          split="train",
          dtype=dtype,
          shuffle_seed=shuffle_seed,
      )
  
      if dataset_service_address:
          if shuffle_seed is not None:
              raise ValueError(
                  "Using dataset service with a random seed causes each "
                  "worker to produce exactly the same data. Add "
                  "config.shuffle_seed = None to your config if you "
                  "want to run with dataset service."
              )
          logging.info(
              "Using the tf.data service at %s", dataset_service_address
          )
          train_ds = dataset_utils.distribute(
              train_ds, dataset_service_address
          )
  
      logging.info("Loading validation split of the ShapeNet dataset.")
      val_ds, _ = shapenet_load_split(
          eval_batch_size,
          shapenet_config,
          split="validation",
          dtype=dtype,
      )
  
      logging.info("Loading test split of the ShapeNet dataset.")
      test_ds, _ = shapenet_load_split(
          eval_batch_size, shapenet_config, split="test", dtype=dtype
      )
  
      maybe_pad_batches_train = functools.partial(
          dataset_utils.maybe_pad_batch,
          train=True,
          batch_size=batch_size,
          pixel_level=True,
      )
      maybe_pad_batches_val = functools.partial(
          dataset_utils.maybe_pad_batch,
          train=False,
          batch_size=eval_batch_size,
          pixel_level=True,
      )
      maybe_pad_batches_test = functools.partial(
          dataset_utils.maybe_pad_batch,
          train=False,
          batch_size=eval_batch_size,
          pixel_level=True,
      )
      shard_batches = functools.partial(
          dataset_utils.shard, n_devices=num_shards
      )
  
      train_iter = iter(train_ds)
      train_iter = map(dataset_utils.tf_to_numpy, train_iter)
      train_iter = map(maybe_pad_batches_train, train_iter)
      train_iter = map(shard_batches, train_iter)
      train_iter = jax_utils.prefetch_to_device(
          train_iter, prefetch_buffer_size
      )
  
      val_iter = iter(val_ds)
      val_iter = map(dataset_utils.tf_to_numpy, val_iter)
      val_iter = map(maybe_pad_batches_val, val_iter)
      val_iter = map(shard_batches, val_iter)
      val_iter = jax_utils.prefetch_to_device(
          val_iter, prefetch_buffer_size
      )
  
      test_iter = iter(test_ds)
      test_iter = map(dataset_utils.tf_to_numpy, test_iter)
      test_iter = map(maybe_pad_batches_test, test_iter)
      test_iter = map(shard_batches, test_iter)
      test_iter = jax_utils.prefetch_to_device(
          test_iter, prefetch_buffer_size
      )
  
      input_shape = (-1, seq_length, 3)
      label_shape = (-1, shapenet_config.num_category_classes)
  
      meta_data = {
          "num_classes": shapenet_config.num_classes,
          "input_shape": input_shape,
          "label_shape": label_shape,
          "num_train_examples": shapenet_config.num_train_pointclouds,
          "num_eval_examples": shapenet_config.num_validation_pointclouds,
          "num_test_examples": shapenet_config.num_test_pointclouds,
          "input_dtype": getattr(jnp, dtype_str),
          "label_dtype": getattr(jnp, "int64"),
          "target_is_onehot": onehot_labels,
      }
      return dataset_utils.Dataset(
          train_iter, val_iter, test_iter, meta_data
      )
  
file_path: scenic/scenic/projects/pointcloud/shapenet_dataset.py
imports: import collections
  import dataclasses
  import functools
  from typing import Optional
  from absl import logging
  from flax import jax_utils
  import jax.numpy
  import numpy
  from scenic.dataset_lib import dataset_utils
  from scenic.dataset_lib import datasets
  import tensorflow
label_desc: Use image node_id 4 for calling a global function with example usage: get_dataset() and returns: dataset_utils