def create_val_dataset(config, split, batch_size, pad_last_batch):
  """Create validataion dataset.

  Args:
    config: ml_collections.ConfigDict to use.
    split: The validation split.
    batch_size: The batch size.
    pad_last_batch: Bool to indicate whether to pad last patch or not.

  Returns:
    The validation dataset.
  """
  dataset_builder = gscan_dataset.GSCANDataset(**config)
  num_batches = None
  cardinality = None
  if pad_last_batch:
    num_examples = dataset_builder.num_examples[split]
    val_batch_size = jax.local_device_count() * batch_size
    num_batches = int(
        np.ceil(num_examples / val_batch_size / jax.process_count()))
    cardinality = int(np.ceil(num_examples / jax.process_count()))
  ds = deterministic_data.create_dataset(
      dataset_builder,
      split=split,
      preprocess_fn=dataset_builder.preprocess,
      cache=jax.process_count() > 1,
      batch_dims=[jax.local_device_count(), batch_size],
      num_epochs=1,
      shuffle=False,
      pad_up_to_batches=num_batches,
      cardinality=cardinality)
  return ds
Create validataion dataset.

  Args:
    config: ml_collections.ConfigDict to use.
    split: The validation split.
    batch_size: The batch size.
    pad_last_batch: Bool to indicate whether to pad last patch or not.

  Returns:
    The validation dataset.
  
language/gscan/xattn_model/dataset/input_pipeline.py