code: def sample_batch(
      key: jnp.ndarray,
      data: jnp.ndarray,
      labels: jnp.ndarray,
      batch_size: int,
  ) -> Tuple:
      ix = random.choice(
          key=key, a=data.shape[0], shape=(batch_size,), replace=False
      )
      return (
          jnp.take(data, indices=ix, axis=0),
          jnp.take(labels, indices=ix, axis=0),
      )
  
file_path: evojax/evojax/task/mnist.py
imports: import sys
  import numpy
  from typing import Tuple
  import jax
  import jax.numpy
  from jax import random
  from flax.struct import dataclass
  from evojax.task.base import VectorizedTask
  from evojax.task.base import TaskState
label_desc: Use image node_id 1 for calling a global function with example usage: sample_batch(key, data, labels, batch_size) and returns: jnp,  jnp