code: def train_step(state, batch, learning_rate_fn):
      """Perform a single training step."""
  
      def loss_fn(params):
          """loss function used for training."""
          logits, new_model_state = state.apply_fn(
              {"params": params, "batch_stats": state.batch_stats},
              batch["image"],
              mutable=["batch_stats"],
          )
          loss = cross_entropy_loss(logits, batch["label"])
          weight_penalty_params = jax.tree_util.tree_leaves(params)
          weight_decay = 0.0001
          weight_l2 = sum(
              jnp.sum(x**2)
              for x in weight_penalty_params
              if x.ndim > 1
          )
          weight_penalty = weight_decay * 0.5 * weight_l2
          loss = loss + weight_penalty
          return loss, (new_model_state, logits)
  
      step = state.step
      dynamic_scale = state.dynamic_scale
      lr = learning_rate_fn(step)
  
      if dynamic_scale:
          grad_fn = dynamic_scale.value_and_grad(
              loss_fn, has_aux=True, axis_name="batch"
          )
          dynamic_scale, is_fin, aux, grads = grad_fn(state.params)
          # dynamic loss takes care of averaging gradients across replicas
      else:
          grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
          aux, grads = grad_fn(state.params)
          # Re-use same axis_name as in the call to `pmap(...train_step...)` below.
          grads = lax.pmean(grads, axis_name="batch")
      new_model_state, logits = aux[1]
      metrics = compute_metrics(logits, batch["label"])
      metrics["learning_rate"] = lr
  
      new_state = state.apply_gradients(
          grads=grads, batch_stats=new_model_state["batch_stats"]
      )
      if dynamic_scale:
          # if is_fin == False the gradients contain Inf/NaNs and optimizer state and
          # params should be restored (= skip this step).
          new_state = new_state.replace(
              opt_state=jax.tree_util.tree_map(
                  functools.partial(jnp.where, is_fin),
                  new_state.opt_state,
                  state.opt_state,
              ),
              params=jax.tree_util.tree_map(
                  functools.partial(jnp.where, is_fin),
                  new_state.params,
                  state.params,
              ),
              dynamic_scale=dynamic_scale,
          )
          metrics["scale"] = dynamic_scale.scale
  
      return new_state, metrics
  
imports: import functools
  import time
  from typing import Any
  from absl import logging
  from clu import metric_writers
  from clu import periodic_actions
  from flax import jax_utils
  from flax.training import checkpoints
  from flax.training import common_utils
  from flax.training import dynamic_scale
  from flax.training import train_state
  import jax
  from jax import lax
  import jax.numpy
  from jax import random
  import ml_collections
  import optax
  import tensorflow
  import tensorflow_datasets
  import input_pipeline
  import models
label_desc: Use image node_id 6 for calling a global function with example usage: train_step(state, batch, learning_rate_fn) and returns: new_state,  metrics, loss,  new_model_state,  logits