code: def _create_loss_metrics(
      loss_has_aux: bool,
      loss_result: Union[
          jnp.ndarray, Tuple[jnp.ndarray, loggers.LoggingData]
      ],
      gradients: jnp.ndarray,
  ):
      """Creates loss metrics for logging."""
      # Validate input.
      if loss_has_aux and not (
          len(loss_result) == 2
          and isinstance(loss_result[0], jnp.ndarray)
          and isinstance(loss_result[1], dict)
      ):
          raise ValueError(
              "Could not parse loss value and metrics from loss_fn's "
              "output. Since loss_has_aux is enabled, loss_fn must "
              "return loss_value and auxiliary metrics."
          )
  
      if not loss_has_aux and not isinstance(loss_result, jnp.ndarray):
          raise ValueError(
              f"Loss returns type {loss_result}. However, it should "
              "return a jnp.ndarray, given that loss_has_aux = False."
          )
  
      # Maybe unpack loss result.
      if loss_has_aux:
          loss, metrics = loss_result
      else:
          loss = loss_result
          metrics = {}
  
      # Complete metrics dict and return it.
      metrics["loss"] = loss
      metrics["gradient_norm"] = optax.global_norm(gradients)
      return metrics
  
imports: import time
  from typing import Dict, List, NamedTuple, Optional, Tuple, Union, Iterator
  import acme
  from acme import types
  from acme.agents.jax.bc import losses
  from acme.agents.jax.bc import networks
  from acme.jax import networks
  from acme.jax import utils
  from acme.utils import counting
  from acme.utils import loggers
  import jax
  import jax.numpy
  import optax
label_desc: Use image node_id 1 for calling a global function with example usage: _create_loss_metrics(loss_has_aux, loss_result, gradients) and returns: metrics