def testReturnsCorrectAnchorWiseLossWithHighLogitScaleSetting(self):
    """At very high logit_scale, all predictions will be ~0.33."""
    def graph_fn():
      # TODO(yonib): Also test logit_scale with anchorwise=False.
      logit_scale = 10e16
      prediction_tensor = tf.constant([[[-100, 100, -100],
                                        [100, -100, -100],
                                        [0, 0, -100],
                                        [-100, -100, 100]],
                                       [[-100, 0, 0],
                                        [-100, 100, -100],
                                        [-100, 100, -100],
                                        [100, -100, -100]]], tf.float32)
      target_tensor = tf.constant([[[0, 1, 0],
                                    [1, 0, 0],
                                    [1, 0, 0],
                                    [0, 0, 1]],
                                   [[0, 0, 1],
                                    [0, 1, 0],
                                    [0, 1, 0],
                                    [1, 0, 0]]], tf.float32)
      weights = tf.constant([[[1, 1, 1],
                              [1, 1, 1],
                              [1, 1, 1],
                              [1, 1, 1]],
                             [[1, 1, 1],
                              [1, 1, 1],
                              [1, 1, 1],
                              [1, 1, 1]]], tf.float32)
      loss_op = losses.WeightedSoftmaxClassificationLoss(
          logit_scale=logit_scale)
      loss = loss_op(prediction_tensor, target_tensor, weights=weights)
      return loss
    uniform_distribution_loss = - math.log(.33333333333)
    exp_loss = np.matrix([[uniform_distribution_loss] * 4,
                          [uniform_distribution_loss] * 4])
    loss_output = self.execute(graph_fn, [])
    self.assertAllClose(loss_output, exp_loss)
At very high logit_scale, all predictions will be ~0.33.
inference_api/src/main/object_detection/core/losses_test.py