code: def load_tfds_dataset(
      dataset_name: str,
      num_episodes: Optional[int] = None,
      env_spec: Optional[specs.EnvironmentSpec] = None,
  ) -> tf.data.Dataset:
      """Returns a TFDS dataset with the given name."""
      # Used only in tests.
      del env_spec
  
      dataset = tfds.load(dataset_name)["train"]
      if num_episodes:
          dataset = dataset.take(num_episodes)
      return dataset
  
imports: import logging
  from typing import Any, Iterator, Optional, Tuple, Sequence
  from acme import specs
  from acme import types
  from flax import jax_utils
  import jax
  import jax.numpy
  import numpy
  import rlds
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 4 for calling a global function with example usage: load_tfds_dataset(dataset_name, num_episodes, env_spec) and returns: dataset