code: def train_step(
      train_state: train_utils.TrainState,
      batch: Batch,
      task: str,
      dataset: str,
      flax_model: nn.Module,
      loss_fn: LossFn,
      lr_fn: LrFn,
      metrics_fn: MetricFn,
      config: ml_collections.ConfigDict,
      debug: Optional[bool] = False,
  ) -> Tuple[
      train_utils.TrainState,
      Dict[str, Tuple[float, int]],
      Dict[str, Any],
  ]:
      """Runs a single step of training.
  
      This function is branched from scenic/train_lib/classification_trainer.py.
      Here, the model function takes two additional args, `task` and `dataset`.
  
      Args:
        train_state: The state of training including the current global_step,
          model_state, rng, params, and optimizer. The buffer of this argument can
          be donated to the computation.
        batch: A single batch of data. The buffer of this argument can be donated to
          the computation.
        task: The task name, 'temporal_localization', 'moment_retrieval',
          'highlight_detection' or 'action_segmentation'.
        dataset: The dataset name.
        flax_model: A Flax model.
        loss_fn: A loss function that given logits, a batch, and parameters of the
          model calculates the loss.
        lr_fn: The learning rate fn used for the logging the learning rate.
        metrics_fn: A metrics function that given logits and batch of data,
          calculates the metrics as well as the loss.
        config: Configurations of the experiment.
        debug: Whether the debug mode is enabled during training. `debug=True`
          enables model specific logging/storing some values using
          jax.host_callback.
  
      Returns:
        Updated state of training and computed metrics and some training logs.
      """
  
      training_logs = {}
      new_rng, rng = jax.random.split(train_state.rng)
  
      if config.get("mixup") and config.mixup.alpha:
          mixup_rng, rng = jax.random.split(rng, 2)
          mixup_rng = train_utils.bind_rng_to_host_device(
              mixup_rng,
              axis_name="batch",
              bind_to=config.mixup.get("bind_to", "device"),
          )
          batch["rgb"] = batch["inputs"]["rgb"]
          batch = dataset_utils.mixup(
              batch,
              config.mixup.alpha,
              config.mixup.get("image_format", "NHWC"),
              input_key="rgb",
              rng=mixup_rng,
          )
          batch["inputs"]["rgb"] = batch.pop("rgb")
  
      # Bind the rng to the host/device we are on.
      dropout_rng = train_utils.bind_rng_to_host_device(
          rng, axis_name="batch", bind_to="device"
      )
  
      all_gather_loss = config.get("all_gather_loss", False)
      gathered_batch = (
          unloc_eval_utils.all_gather_metrics_inputs(batch)
          if all_gather_loss
          else None
      )
  
      def training_loss_fn(params):
          variables = {"params": params, **train_state.model_state}
          logits = flax_model.apply(
              variables,
              batch["inputs"],
              task=task,
              dataset=dataset,
              train=True,
              rngs={"dropout": dropout_rng},
              debug=debug,
          )
          loss = loss_fn(logits, batch, variables["params"])
          return loss, logits
  
      def training_loss_fn_all_gather(params):
          variables = {"params": params, **train_state.model_state}
          logits = unloc_eval_utils.run_model_all_gather_results(
              variables,
              batch,
              task,
              flax_model,
              train=True,
              dropout_rng=dropout_rng,
              debug=debug,
          )
          loss = loss_fn(logits, gathered_batch, params)
          return loss, logits
  
      compute_gradient_fn = jax.value_and_grad(
          training_loss_fn_all_gather
          if all_gather_loss
          else training_loss_fn,
          has_aux=True,
      )
      (train_cost, logits), grad = compute_gradient_fn(
          train_state.params
      )
  
      del train_cost
      # Re-use same axis_name as in the call to `pmap(...train_step...)` below.
      grad = jax.lax.pmean(grad, axis_name="batch")
  
      if config.get("max_grad_norm") is not None:
          grad = clip_grads(grad, config.max_grad_norm)
  
      assert train_state.tx is not None
      updates, new_opt_state = train_state.tx.update(
          grad, train_state.opt_state, train_state.params
      )
      new_params = optax.apply_updates(train_state.params, updates)
  
      training_logs["l2_grads"] = jnp.sqrt(
          sum([jnp.vdot(g, g) for g in jax.tree_util.tree_leaves(grad)])
      )
      ps = jax.tree_util.tree_leaves(new_params)
      training_logs["l2_params"] = jnp.sqrt(
          sum([jnp.vdot(p, p) for p in ps])
      )
      us = jax.tree_util.tree_leaves(updates)
      training_logs["l2_updates"] = jnp.sqrt(
          sum([jnp.vdot(u, u) for u in us])
      )
      # TODO(dehghani): Can we get this from the optimizer instead?
      training_logs["learning_rate"] = lr_fn(train_state.global_step)
  
      metrics = metrics_fn(
          logits, gathered_batch if all_gather_loss else batch
      )
  
      new_train_state = (
          train_state.replace(  # pytype: disable=attribute-error
              global_step=train_state.global_step + 1,
              opt_state=new_opt_state,
              params=new_params,
              rng=new_rng,
          )
      )
  
      return new_train_state, metrics, training_logs
  
imports: import collections.abc
  import functools
  import re
  from typing import Any, Callable, Dict, Mapping, Optional, Sequence, Tuple, Union
  from absl import logging
  import flax
  import flax.linen
  import jax
  from jax.example_libraries.optimizers import clip_grads
  import jax.numpy
  import ml_collections
  import optax
  from scenic.common_lib import debug_utils
  from scenic.dataset_lib import dataset_utils
  from scenic.projects.baselines.clip import model
  from scenic.projects.unloc import eval_utils
  from scenic.projects.unloc import model_utils
  from scenic.projects.vivit import model_utils
  from scenic.train_lib import pretrain_utils
  from scenic.train_lib import train_utils
  import scipy.ndimage
label_desc: Use image node_id 10 for calling a global function with example usage: train_step(train_state, batch, task, dataset, flax_model, loss_fn, lr_fn, metrics_fn, config, debug) and returns: new_train_state,  metrics,  training_logs, loss,  logits, loss,  logits