class_docstr: An abstract class that defines a data source.
class_docstr_tok: ["An","abstract","class","that","defines","a","data","source","."]
docstr: Crates a dataset completely preloaded in memory.
  
  This creates a tf.data.Dataset which is constructed by load all data
  into memory and pre-shuffling (if applicable). This is much faster than
  having tf.data.Dataset handle individual items.
  
  Args:
    item_ids: the item IDs to construct the datset with.
    flatten: whether to flatten the image dimensions.
    shuffle: whether to shuffle the dataset.
  
  Returns:
    A tf.data.Dataset instance.
code: def _create_preloaded_dataset(
      self, item_ids, flatten=False, shuffle=False
  ):
      """Crates a dataset completely preloaded in memory.
  
      This creates a tf.data.Dataset which is constructed by load all data
      into memory and pre-shuffling (if applicable). This is much faster than
      having tf.data.Dataset handle individual items.
  
      Args:
        item_ids: the item IDs to construct the datset with.
        flatten: whether to flatten the image dimensions.
        shuffle: whether to shuffle the dataset.
  
      Returns:
        A tf.data.Dataset instance.
      """
      load_fn = functools.partial(self.get_item)
      data_list = utils.parallel_map(load_fn, item_ids)
      data_list = [_camera_to_rays_fn(item) for item in data_list]
      data_dict = utils.tree_collate(data_list)
  
      num_examples = data_dict["origins"].shape[0]
      heights = [x.shape[0] for x in data_dict["origins"]]
      widths = [x.shape[1] for x in data_dict["origins"]]
  
      # Broadcast appearance ID to match ray shapes.
      if "metadata" in data_dict:
          for metadata_key, metadata in data_dict["metadata"].items():
              data_dict["metadata"][metadata_key] = np.asarray(
                  [
                      np.full((heights[i], widths[i], 1), fill_value=x)
                      for i, x in enumerate(metadata)
                  ]
              )
  
      num_rays = int(sum([x * y for x, y in zip(heights, widths)]))
      shuffled_inds = self.rng.permutation(num_rays)
  
      logging.info(
          "*** Loaded dataset items: num_rays=%d, num_examples=%d",
          num_rays,
          num_examples,
      )
  
      def _prepare_array(x):
          if not isinstance(x, np.ndarray):
              x = np.asarray(x)
          # Create last dimension if it doesn't exist.
          # The `and` part of the check ensures we're not touching ragged arrays.
          if x.ndim == 1 and x[0].ndim == 0:
              x = np.expand_dims(x, -1)
          if flatten:
              x = np.concatenate(
                  [x.reshape(-1, x.shape[-1]) for x in x], axis=0
              )
          if shuffle:
              x = x[shuffled_inds]
          return x
  
      out_dict = {}
      for key, value in data_dict.items():
          out_dict[key] = jax.tree_map(_prepare_array, value)
  
      return tf.data.Dataset.from_tensor_slices(out_dict)
  
imports: import abc
  import functools
  import itertools
  from typing import Any, Iterable, Optional, Sequence, Union
  from absl import logging
  from flax import jax_utils
  import jax
  import numpy
  import tensorflow
  from nerfies import camera
  from nerfies import gpath
  from nerfies import image_utils
  from nerfies import tf_camera
  from nerfies import utils
  from tensorflow.python.data.util import nest
label_desc: Use image node_id 26 for calling the DataSource obj's underlying member method code with example usage: obj._create_preloaded_dataset(item_ids, flatten, shuffle) and returns: tf, x