code: def get_distillation_loss_fn(
      teacher_logits, distillation_spec, my_id, original_loss_fn
  ):
      """Force the loss_fn to compare the student logits to the teacher logits."""
  
      # The logits input below is the student logits
      def _mse_teacher_loss_fn(labels, logits, weights=1.0):
          """A mean square error with the teacher's logits/predictions."""
          del labels  # Unused.
  
          return tf.compat.v1.losses.mean_squared_error(
              labels=teacher_logits, predictions=logits, weights=weights
          )
  
      def _cross_entropy_loss_fn(labels, logits, weights=1.0):
          """A cross entropy with the teacher's predictions."""
          del labels  # Unused.
  
          return tf.compat.v1.losses.softmax_cross_entropy(
              onehot_labels=teacher_logits,
              logits=logits,
              weights=weights,
          )
  
      # TODO(b/172564129): Once b/148613464 is in place, add * to force parameter
      # names.
      def _adaptively_balance_losses_loss_fn(
          loss1_fn, loss2_fn, labels, logits, weights=1.0
      ):
          """Increasingly grow the distillation loss over time."""
          original_loss = loss1_fn(
              labels=labels, logits=logits, weights=weights
          )
          distill_loss = loss2_fn(
              labels=labels, logits=logits, weights=weights
          )
  
          lambda_ = distillation_spec.balance_losses_spec.minimum_lambda
          delta = (
              distillation_spec.balance_losses_spec.maximum_lambda
              - distillation_spec.balance_losses_spec.minimum_lambda
          )
          max_trials = (
              distillation_spec.balance_losses_spec.num_ramp_trials
          )
  
          if max_trials > 0:
              # Linearly increase the lambda toward the maximum
              lambda_ += delta * (
                  max(0, my_id - distillation_spec.minimal_pool_size)
                  / max_trials
              )
  
          return original_loss + lambda_ * distill_loss
  
      if (
          distillation_spec.distillation_type
          == distillation_spec_pb2.DistillationSpec.MSE_LOGITS
          or distillation_spec.distillation_type
          == distillation_spec_pb2.DistillationSpec.MSE_SOFTMAX
      ):
          return _mse_teacher_loss_fn
      elif (
          distillation_spec.distillation_type
          == distillation_spec_pb2.DistillationSpec.CROSS_ENTROPY
      ):
          return _cross_entropy_loss_fn
      elif (
          distillation_spec.distillation_type
          == distillation_spec_pb2.DistillationSpec.ADAPTIVELY_BALANCE_LOSSES
      ):
          return functools.partial(
              _adaptively_balance_losses_loss_fn,
              loss1_fn=original_loss_fn,
              loss2_fn=_cross_entropy_loss_fn,
          )
  
      else:
          raise ValueError("unsupported distillation_type.")
  
file_path: model_search/model_search/meta/distillation.py
imports: import collections
  import functools
  from model_search import logit_bundler
  from model_search.architecture import architecture_utils
  from model_search.proto import distillation_spec_pb2
  import tensorflow.compat.v2
label_desc: Use image node_id 1 for calling a global function with example usage: get_distillation_loss_fn(teacher_logits, distillation_spec, my_id, original_loss_fn) and returns: tf, tf, unknown, _mse_teacher_loss_fn, _cross_entropy_loss_fn, functools