Use tf.squeeze before casting non-0d tensors to float (#370)

This becomes a hard error after numpy 2.4 which we are now using
internally.

Author: boomanaiden154

gematria/model/python/loss_utils_test.py

@@ -66,10 +66,14 @@ def test_unscaled_loss(self):
     mse = loss.mean_squared_error
     mae = loss.mean_absolute_error
     percentiles = loss.absolute_error_percentiles
-    self.assertNear(float(mse), (3**2 + 0 + 0 + 4**2 + 0.5**2) / 5, 1e-6)
-    self.assertNear(float(mae), (3 + 0 + 0 + 4 + 0.5) / 5, 1e-6)
     self.assertNear(
-        float(loss.loss_tensor), (2.5 + 0 + 0 + 3.5 + (0.5**2) / 2) / 5, 1e-6
+        float(tf.squeeze(mse)), (3**2 + 0 + 0 + 4**2 + 0.5**2) / 5, 1e-6
+    )
+    self.assertNear(float(tf.squeeze(mae)), (3 + 0 + 0 + 4 + 0.5) / 5, 1e-6)
+    self.assertNear(
+        float(tf.squeeze(loss.loss_tensor)),
+        (2.5 + 0 + 0 + 3.5 + (0.5**2) / 2) / 5,
+        1e-6,
     )
     self.assertAllEqual(percentiles, ((0,), (0,), (0.5,), (3,), (4,)))
 
@@ -86,9 +90,12 @@ def test_percentage_loss(self):
     mape = loss.mean_absolute_percentage_error
     percentiles = loss.absolute_percentage_error_percentiles
     self.assertAlmostEqual(
-        float(mspe), ((3 / 4) ** 2 + 0 + 0 + 4**2 + (0.5 / 2) ** 2) / 5
+        float(tf.squeeze(mspe)),
+        ((3 / 4) ** 2 + 0 + 0 + 4**2 + (0.5 / 2) ** 2) / 5,
+    )
+    self.assertAlmostEqual(
+        float(tf.squeeze(mape)), (3 / 4 + 0 + 0 + 4 + 0.5 / 2) / 5
     )
-    self.assertAlmostEqual(float(mape), (3 / 4 + 0 + 0 + 4 + 0.5 / 2) / 5)
     self.assertAllEqual(percentiles, ((0,), (0.5 / 2,), (0.5 / 2,), (4,)))
 
   def test_normalized_loss_when_expected_value_greater_than_one(self):
@@ -121,11 +128,11 @@ def test_normalized_loss_when_expected_value_greater_than_one(self):
         loss_type=options.LossType.MEAN_SQUARED_ERROR
     )
     self.assertAlmostEqual(
-        float(mean_absolute_error.loss_tensor),
+        float(tf.squeeze(mean_absolute_error.loss_tensor)),
         (0.3 + 1.5 + 0.0 + 0.5 + 2.0) / 5,
     )
     self.assertAlmostEqual(
-        float(mean_squared_error.loss_tensor),
+        float(tf.squeeze(mean_squared_error.loss_tensor)),
         (0.3**2 + 1.5**2 + 0.0 + 0.5**2 + 2.0**2) / 5,
         delta=1e-6,
     )
@@ -272,10 +279,14 @@ def test_unknown_shape(self):
     self.assertEqual(loss.loss_tensor.shape, (1,))
     self.assertEqual(percentiles.shape, (len(percentile_ranks), 1))
 
-    self.assertNear(float(mse), (3**2 + 0 + 0 + 4**2 + 0.5**2) / 5, 1e-6)
-    self.assertNear(float(mae), (3 + 0 + 0 + 4 + 0.5) / 5, 1e-6)
     self.assertNear(
-        float(loss.loss_tensor), (2.5 + 0 + 0 + 3.5 + (0.5**2) / 2) / 5, 1e-6
+        float(tf.squeeze(mse)), (3**2 + 0 + 0 + 4**2 + 0.5**2) / 5, 1e-6
+    )
+    self.assertNear(float(tf.squeeze(mae)), (3 + 0 + 0 + 4 + 0.5) / 5, 1e-6)
+    self.assertNear(
+        float(tf.squeeze(loss.loss_tensor)),
+        (2.5 + 0 + 0 + 3.5 + (0.5**2) / 2) / 5,
+        1e-6,
     )
     self.assertAllEqual(percentiles, ((0,), (0.5,), (3,), (4,)))
 
@@ -349,10 +360,14 @@ def test_single_task_unknown_shape(self):
     self.assertEqual(loss.loss_tensor.shape, (num_tasks,))
     self.assertEqual(percentiles.shape, (len(percentile_ranks), num_tasks))
 
-    self.assertNear(float(mse), (3**2 + 0 + 0 + 4**2 + 0.5**2) / 5, 1e-6)
-    self.assertNear(float(mae), (3 + 0 + 0 + 4 + 0.5) / 5, 1e-6)
     self.assertNear(
-        float(loss.loss_tensor), (2.5 + 0 + 0 + 3.5 + (0.5**2) / 2) / 5, 1e-6
+        float(tf.squeeze(mse)), (3**2 + 0 + 0 + 4**2 + 0.5**2) / 5, 1e-6
+    )
+    self.assertNear(float(tf.squeeze(mae)), (3 + 0 + 0 + 4 + 0.5) / 5, 1e-6)
+    self.assertNear(
+        float(tf.squeeze(loss.loss_tensor)),
+        (2.5 + 0 + 0 + 3.5 + (0.5**2) / 2) / 5,
+        1e-6,
     )
     self.assertAllEqual(percentiles, ((0,), (0.5,), (3,), (4,)))
 

gematria/model/python/model_base_test.py

@@ -359,7 +359,7 @@ def test_randomized_expected_outputs_delta(self):
       for i, prefix_throughputs in enumerate(
           inverse_throughputs.prefix_inverse_throughput_cycles
       ):
-        expected_output_delta = float(expected_output_deltas[i])
+        expected_output_delta = float(tf.squeeze(expected_output_deltas[i]))
         possible_deltas = set()
         for prefix_throughput in prefix_throughputs:
           for previous_prefix_throughput in previous_prefix_throughputs:
@@ -676,10 +676,10 @@ def test_training_with_full_variable_list(self):
     )
     biases = model._variable_groups[TestModelWithVarGroups.BIAS]
     for bias in biases:
-      self.assertNotAlmostEqual(float(bias), -0.5)
+      self.assertNotAlmostEqual(float(tf.squeeze(bias)), -0.5)
     weights = model._variable_groups[TestModelWithVarGroups.WEIGHTS]
     for weight in weights:
-      self.assertNotAlmostEqual(float(weight), 0.5)
+      self.assertNotAlmostEqual(float(tf.squeeze(weight)), 0.5)
 
   def test_training_bias_only(self):
     task_list = ['foo', 'bar']
@@ -698,10 +698,10 @@ def test_training_bias_only(self):
     )
     biases = model._variable_groups[TestModelWithVarGroups.BIAS]
     for bias in biases:
-      self.assertNotAlmostEqual(float(bias), -0.5)
+      self.assertNotAlmostEqual(float(tf.squeeze(bias)), -0.5)
     weights = model._variable_groups[TestModelWithVarGroups.WEIGHTS]
     for weight in weights:
-      self.assertAlmostEqual(float(weight), 0.5)
+      self.assertAlmostEqual(float(tf.squeeze(weight)), 0.5)
 
   def test_grad_clipping(self):
     task_list = ['foo', 'bar']
@@ -737,10 +737,10 @@ def test_training_weight_only(self):
     )
     biases = model._variable_groups[TestModelWithVarGroups.BIAS]
     for bias in biases:
-      self.assertAlmostEqual(float(bias), -0.5)
+      self.assertAlmostEqual(float(tf.squeeze(bias)), -0.5)
     weights = model._variable_groups[TestModelWithVarGroups.WEIGHTS]
     for weight in weights:
-      self.assertNotAlmostEqual(float(weight), 0.5)
+      self.assertNotAlmostEqual(float(tf.squeeze(weight)), 0.5)
 
 
 if __name__ == '__main__':