code: def modelnet40_load_split(
      batch_size,
      dataset_config,
      train,
      onehot_labels=True,
      dtype=tf.float32,
      prefetch_buffer_size=10,
      shuffle_seed=None,
  ):
      """Creates a split from the ModelNet40 dataset using TensorFlow Datasets.
  
      For the training set, we drop the last partial batch. This is fine to do
      because we additionally shuffle the data randomly each epoch, thus the trainer
      will see all data in expectation. For the validation set, we pad the final
      batch to the desired batch size.
  
      Args:
        batch_size: int; The batch size returned by the data pipeline.
        dataset_config: dataset configuration.
        train: bool; Whether to load the train or evaluation split.
        onehot_labels: Whether to transform the labels to one hot.
        dtype: TF data type; Data type of the image.
        prefetch_buffer_size: int; Buffer size for the TFDS prefetch.
        shuffle_seed: The seed to use when shuffling the train split.
  
      Returns:
        A `tf.data.Dataset`.
      """
  
      def random_noise(pointcloud, std=0.02):
          assert len(pointcloud.shape) == 2
          noise = tf.random.normal(
              tf.shape(pointcloud), mean=0, stddev=std / 2
          )
          noisy_pointcloud = pointcloud + tf.clip_by_value(
              noise, clip_value_min=-std, clip_value_max=std
          )
          return noisy_pointcloud
  
      def decode_example(example):
          pointcloud = tf.cast(example["pc"], dtype=dtype)
          if train:
              pointcloud = random_noise(pointcloud)
              pointcloud = tf.random.shuffle(pointcloud)
  
          label = example["label"]
          label = (
              tf.one_hot(label, dataset_config.num_classes)
              if onehot_labels
              else label
          )
          return {"inputs": pointcloud, "label": label}
  
      split = "train" if train else "test"
      dataset_builder = dataset_utils.get_dataset_tfds(
          dataset_config.dataset_name,
          split=split,
          data_dir=dataset_config.dataset_path,
          skip_decode=("pc",),
      )
      # Download dataset:
      dataset_builder.download_and_prepare()
  
      ds = dataset_builder.as_dataset(
          split=split,
          decoders={
              "pc": tfds.decode.SkipDecoding(),
          },
      )
      options = tf.data.Options()
      options.threading.private_threadpool_size = 48
      ds = ds.with_options(options)
  
      ds = ds.cache()
  
      if train:
          ds = ds.repeat()
          ds = ds.shuffle(16 * batch_size, seed=shuffle_seed)
  
      # decode_example should be applied after caching as it also does augmentation
      ds = ds.map(
          decode_example,
          num_parallel_calls=tf.data.experimental.AUTOTUNE,
      )
      ds = ds.batch(batch_size, drop_remainder=train)
  
      ds = ds.prefetch(prefetch_buffer_size)
      return ds
  
file_path: scenic/scenic/projects/pointcloud/pointcloud_dataset.py
imports: import dataclasses
  import functools
  from typing import Optional
  from absl import logging
  from flax import jax_utils
  import jax.numpy
  from scenic.dataset_lib import dataset_utils
  from scenic.dataset_lib import datasets
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 2 for calling a global function with example usage: modelnet40_load_split(batch_size, dataset_config, train, onehot_labels, dtype, prefetch_buffer_size, shuffle_seed) and returns: ds, noisy_pointcloud, dict