text: from elegy.module import Defered
  from typing import Tuple
  from elegy.metrics.metric import Metric
  import typing as tp
  import haiku as hk
  
  from elegy import utils
  import jax
  
  
  def forward_all(metrics):
      def _metrics_fn(**kwargs):
  
          logs = {}
  
          for context, val in apply_recursive((), metrics, **kwargs):
              name = "/".join(context)
              loss_name = get_unique_name(logs, name)
              logs[loss_name] = val
  
          return logs
  
      return _metrics_fn
  
  
  def apply_recursive(context: tp.Tuple[str, ...], metrics, **kwargs):
  
      if isinstance(metrics, tp.Callable):
  
          if isinstance(metrics, Defered):
              metrics = metrics.get_instance()
  
          name = (
              metrics.module_name if isinstance(metrics, hk.Module) else metrics.__name__
          )
          context += (name,)
          loss_val = utils.inject_dependencies(metrics)(**kwargs)
  
          if isinstance(loss_val, tp.Dict):
              for name, loss_val in loss_val.items():
                  yield context + (name,), loss_val
          else:
              yield context, loss_val
  
      elif isinstance(metrics, (tp.Tuple, tp.List)):
          for loss in metrics:
              yield from apply_recursive(context, loss, **kwargs)
      elif isinstance(metrics, tp.Dict):
          for name, loss in metrics.items():
              yield from apply_recursive(context + (name,), loss, **kwargs)
      else:
          raise TypeError(f"Invalid type {type(metrics)}")
  
  
  def get_unique_name(logs, name):
  
      if name not in logs:
          return name
  
      i = 1
      while f"{name}_{i}" in logs:
          i += 1
  
      return f"{name}_{i}"
  