code: def _loss_fn(labels, logits, weights=1.0):
      """Cross entropy loss fn."""
  
      label_ids = tf.squeeze(labels)
  
      if label_ids.dtype == tf.float32:
          label_ids = tf.cast(label_ids, "int32")
      one_hot_labels = tf.one_hot(
          indices=label_ids, depth=logits.shape[-1]
      )
  
      return tf.reduce_mean(
          input_tensor=tf.compat.v1.losses.softmax_cross_entropy(
              onehot_labels=one_hot_labels,
              logits=logits,
              weights=weights,
          )
      )
  
imports: import collections
  from absl.testing import parameterized
  from model_search import hparam
  from model_search import loss_fns
  from model_search import task_manager
  from model_search.architecture import architecture_utils
  from model_search.generators import trial_utils
  from model_search.proto import phoenix_spec_pb2
  import numpy
  import tensorflow.compat.v2
  from google.protobuf import text_format
label_desc: Use image node_id 1 for calling a global function with example usage: _loss_fn(labels, logits, weights) and returns: tf