code: def pmapped_steps(model, config):
      """Returns the pmapped train and eval steps."""
      # Learning rate scheduler.
      learning_rate_fn = lr_schedules.get_learning_rate_fn(config)
      train_step_pmapped = jax.pmap(
          functools.partial(
              train_step,
              flax_model=model.flax_model,
              learning_rate_fn=learning_rate_fn,
              loss_fn=model.loss_function,
              metrics_fn=model.get_metrics_fn("train"),
              config=config,
              debug=config.debug_train,
          ),
          axis_name="batch",
          # We can donate both buffers of train_state and train_batch.
          donate_argnums=(1, 2),
          static_broadcasted_argnums=(0,),  # dataset arg.
      )
      eval_step_pmapped = jax.pmap(
          functools.partial(
              eval_step,
              model=model,
              metrics_fn=model.get_metrics_fn("validation"),
              config=config,
              debug=config.debug_eval,
          ),
          axis_name="batch",
      )
      return train_step_pmapped, eval_step_pmapped
  
imports: import copy
  import dataclasses
  import functools
  import json
  import os
  import random
  from typing import Any, Callable, Dict, List, Optional, Tuple
  from absl import logging
  from clu import metric_writers
  from clu import periodic_actions
  from dmvr import tokenizers
  from flax import jax_utils
  from flax.core import unfreeze
  import flax.linen
  import jax
  import jax.example_libraries.optimizers
  import jax.numpy
  import jax.profiler
  import ml_collections
  import numpy
  from scenic.dataset_lib import dataset_utils
  from scenic.model_lib.base_models import base_model
  from scenic.projects.t5 import model
  from scenic.projects.t5 import tokenizer
  from scenic.projects.vid2seq import load_utils
  from scenic.projects.vid2seq import models
  from scenic.projects.vid2seq import train_utils
  from scenic.projects.vid2seq import dvc_eval
  from scenic.train_lib_deprecated import lr_schedules
  from scenic.train_lib_deprecated import optimizers
  from scenic.train_lib_deprecated import pretrain_utils
  from scenic.train_lib_deprecated import train_utils
label_desc: Use image node_id 6 for calling a global function with example usage: pmapped_steps(model, config) and returns: train_step_pmapped,  eval_step_pmapped