class_docstr: Generates TaskSpecs for various tasks in Phoenix.
docstr: Gets the applicable loss_fn to use.
  
  If head is not None, wraps the head's loss function to match the interface
  Phoenix expects (see loss_fns.py), unless distillation is being used.
  
  Args:
    original_loss_fn: The original loss_fn from the user.
    features: The features pass to model_fn.
    mode: The mode passed to model_fn.
    my_id: My trial id (integer).
    teacher_logits_spec: Logits of the teacher network to use when distilling.
  
  Returns:
    The loss_fn to use.
  
  Raises:
    RuntimeError if unable to find the loss function in the head object.
docstr_tok: ["Gets","the","applicable","loss_fn","to","use",".","If","head","is","not","None",",","wraps","the","head","'s","loss","function","to","match","the","interface","Phoenix","expects","(","see","loss_fns.py",")",",","unless","distillation","is","being","used",".","Args",":","original_loss_fn",":","The","original","loss_fn","from","the","user",".","features",":","The","features","pass","to","model_fn",".","mode",":","The","mode","passed","to","model_fn",".","my_id",":","My","trial","id","(","integer",")",".","teacher_logits_spec",":","Logits","of","the","teacher","network","to","use","when","distilling",".","Returns",":","The","loss_fn","to","use",".","Raises",":","RuntimeError","if","unable","to","find","the","loss","function","in","the","head","object","."]
code: def _get_loss_fn(
      self,
      original_loss_fn,
      features,
      mode,
      my_id,
      teacher_logits_spec=None,
  ):
      """Gets the applicable loss_fn to use.
  
      If head is not None, wraps the head's loss function to match the interface
      Phoenix expects (see loss_fns.py), unless distillation is being used.
  
      Args:
        original_loss_fn: The original loss_fn from the user.
        features: The features pass to model_fn.
        mode: The mode passed to model_fn.
        my_id: My trial id (integer).
        teacher_logits_spec: Logits of the teacher network to use when distilling.
  
      Returns:
        The loss_fn to use.
  
      Raises:
        RuntimeError if unable to find the loss function in the head object.
      """
  
      if (
          mode == tf.estimator.ModeKeys.TRAIN
          and teacher_logits_spec is not None
      ):
          return distillation.get_distillation_loss_fn(
              teacher_logits=teacher_logits_spec.logits,
              distillation_spec=self._phoenix_spec.distillation_spec,
              my_id=my_id,
              original_loss_fn=original_loss_fn,
          )
  
      if not self._head:
          return original_loss_fn
  
      def head_loss_fn(labels, logits, weights=1.0):
          """Create a loss fn from the Head object."""
          del weights  # Head already has weights built in.
  
          training_loss = None
          # There is two types of head, and their api is different.
          if getattr(self._head, "loss", None) is not None:
              training_loss = self._head.loss(
                  labels=labels,
                  logits=logits,
                  features=features,
                  mode=mode,
              )
          elif getattr(self._head, "create_loss", None) is not None:
              training_loss = self._head.create_loss(
                  labels=labels,
                  logits=logits,
                  features=features,
                  mode=mode,
              ).training_loss
          else:
              raise RuntimeError(
                  "unable to find loss function in Head object."
              )
  
          return training_loss
  
      return head_loss_fn
  
imports: import collections
  import functools
  import inspect
  from absl import logging
  from model_search import block_builder
  from model_search import ensembler
  from model_search.architecture import architecture_utils
  from model_search.architecture import tower
  from model_search.ensembler import EnsembleLogits
  from model_search.generators import base_tower_generator
  from model_search.generators import trial_utils
  from model_search.meta import distillation
  import numpy
  import tensorflow.compat.v2
label_desc: Use image node_id 3 for calling the TaskManager obj's underlying member method code with example usage: obj._get_loss_fn(original_loss_fn, features, mode, my_id, teacher_logits_spec) and returns: head_loss_fn, distillation, original_loss_fn, training_loss