Using NumPy APIs to improve performance and the code quality,

tensorflow_ranking/python/losses_test.py

@@ -19,6 +19,8 @@
 from __future__ import print_function
 
 import math
+
+import numpy as np
 import tensorflow as tf
 
 from tensorflow_ranking.python import losses as ranking_losses
@@ -151,14 +153,11 @@ def _loss(si, sj, label_diff, delta):
 
 def _batch_aggregation(batch_loss_list, reduction=None):
   """Returns the aggregated loss."""
-  loss_sum = 0.
-  weight_sum = 0.
-  for loss, weight, count in batch_loss_list:
-    loss_sum += loss
-    if reduction == 'mean':
-      weight_sum += weight
-    else:
-      weight_sum += count
+  loss_sum = np.sum([loss for loss, weight, count in batch_loss_list])
+  if reduction == 'mean':
+    weight_sum = np.sum([weight for loss, weight, count in batch_loss_list])
+  else:
+    weight_sum = np.sum([count for loss, weight, count in batch_loss_list])
   return loss_sum / weight_sum