def create_datasets(config, data_rng):
  """Create datasets for training and evaluation.

  Args:
    config: ml_collections.ConfigDict to use.
    data_rng: PRNGKey for seeding operations in the training dataset.

  Returns:
    A tuple of the training dataset and the evaluation dataset.
  """
  dataset_builder = gscan_dataset.GSCANDataset(**config)
  train_ds = deterministic_data.create_dataset(
      dataset_builder,
      split=config.train_split,
      rng=data_rng,
      preprocess_fn=dataset_builder.preprocess,
      cache=False,
      shuffle_buffer_size=config.shuffle_buffer_size,
      batch_dims=[jax.local_device_count(), config.train_per_device_batch_size],
      num_epochs=config.num_epochs,
      shuffle=True,
      prefetch_size=tf.data.experimental.AUTOTUNE,
  )

  eval_ds = create_val_dataset(config, config.eval_split,
                               config.eval_per_device_batch_size,
                               config.eval_pad_last_batch)
  return train_ds, eval_ds
Create datasets for training and evaluation.

  Args:
    config: ml_collections.ConfigDict to use.
    data_rng: PRNGKey for seeding operations in the training dataset.

  Returns:
    A tuple of the training dataset and the evaluation dataset.
  
language/gscan/xattn_model/dataset/input_pipeline.py