code: def get_datasets() -> tuple[Tree[NumpyElem], Tree[NumpyElem]]:
      """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
  
imports: from __future__ import annotations
  import argparse
  import os
  import typing
  from typing import TYPE_CHECKING
  import attrs
  import cattrs
  import jax
  import jax.numpy
  import numpy
  import optax
  import tensorflow_datasets
  from flax import linen
  from flax import serialization
  from flax.metrics import tensorboard
  from flax.training import train_state
  import bentoml
  from bentoml._internal.utils.pkg import pkg_version_info
label_desc: Use image node_id 4 for calling a global function with example usage: get_datasets() and returns: train_ds,  test_ds