code: def get_tfds_dataset(
      dataset_name: str,
      num_episodes: Optional[int] = None,
      env_spec: Optional[specs.EnvironmentSpec] = None,
  ) -> tf.data.Dataset:
      """Returns a TFDS dataset transformed to a dataset of transitions."""
      dataset = load_tfds_dataset(dataset_name, num_episodes, env_spec)
      batched_steps = dataset.flat_map(_batch_steps)
      return rlds.transformations.map_steps(
          batched_steps, _batched_step_to_transition
      )
  
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 5 for calling a global function with example usage: get_tfds_dataset(dataset_name, num_episodes, env_spec) and returns: rlds