code: def pmapped_steps(model, config):
      """Returns the pmapped train and eval steps."""
      train_step_pmapped = jax.pmap(
          functools.partial(
              train_step,
              flax_model=model.flax_model,
              config=config,
              debug=config.debug_train,
          ),
          axis_name="batch",
          # We can donate both buffers of train_state and train_batch.
          donate_argnums=(0, 1),
      )
      eval_step_pmapped = jax.pmap(
          functools.partial(
              eval_step,
              flax_model=model.flax_model,
              debug=config.debug_eval,
          ),
          axis_name="batch",
          # We can donate the eval_batch's buffer.
          donate_argnums=(1,),
      )
      return train_step_pmapped, eval_step_pmapped
  
imports: import functools
  from typing import Any, Dict, Optional, Tuple
  from absl import logging
  from clu import metric_writers
  from clu import periodic_actions
  from clu import platform
  from flax import jax_utils
  import flax.linen
  import jax
  import jax.example_libraries.optimizers
  import jax.profiler
  import ml_collections
  import numpy
  import optax
  from scenic.projects.verbs_in_action import utils
  from scenic.train_lib import train_utils
label_desc: Use image node_id 4 for calling a global function with example usage: pmapped_steps(model, config) and returns: train_step_pmapped,  eval_step_pmapped