def _model_with_ends(model, end_layers, batch_signature):
  """Returns a model+ends layer built on an already initialized model.

  Ends can be loss or metric layers.

  Args:
    model: Layer with initialized weights and state.
    end_layers: List of end layers.
    batch_signature: Signature of the model input batch.

  Returns:
    An initialized, combined model+ends layer, preserving the initialization
    of ``model``.
  """
  # TODO(jonni): Redo this function as part of an initialization refactor?
  metrics_layer = tl.Branch(*end_layers)
  metrics_input_signature = model.output_signature(batch_signature)
  _, _ = metrics_layer.init(metrics_input_signature)

  model_with_metrics = tl.Serial(model, metrics_layer)
  return model_with_metrics
Returns a model+ends layer built on an already initialized model.

  Ends can be loss or metric layers.

  Args:
    model: Layer with initialized weights and state.
    end_layers: List of end layers.
    batch_signature: Signature of the model input batch.

  Returns:
    An initialized, combined model+ends layer, preserving the initialization
    of ``model``.
  
trax/supervised/training.py