code: def get_datasets():
      """Load MNIST train and test datasets into memory."""
      ds_builder = tfds.builder("mnist")
      ds_builder.download_and_prepare()
      train_ds = tfds.as_numpy(
          ds_builder.as_dataset(split="train", batch_size=-1)
      )
      test_ds = tfds.as_numpy(
          ds_builder.as_dataset(split="test", batch_size=-1)
      )
      train_ds["image"] = jnp.float32(train_ds["image"]) / 255.0
      test_ds["image"] = jnp.float32(test_ds["image"]) / 255.0
      return train_ds, test_ds
  
file_path: flax/examples/mnist/train.py
imports: from absl import logging
  from flax import linen
  from flax.metrics import tensorboard
  from flax.training import train_state
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  import optax
  import tensorflow_datasets
label_desc: Use image node_id 4 for calling a global function with example usage: get_datasets() and returns: train_ds,  test_ds