Use assertAllClose to compare float32 arrays

PiperOrigin-RevId: 723676676
google · Feb 5, 2025 · ccd25a1 · ccd25a1
1 parent ccd78af
commit ccd25a1
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/praxis/sample_decode_test.py b/praxis/sample_decode_test.py
@@ -506,10 +506,10 @@ def decode_fn(model, input_ids, input_paddings):
       # batch size is 1.
       self.assertEqual(1, top_candidate_logprobs.shape[0])
       self.assertEqual(1, top_candidate_ids.shape[0])
-      self.assertArraysEqual(
+      self.assertAllClose(
           logprobs, top_candidate_logprobs[0, :, :, :num_per_token_logprobs]
       )
-      self.assertArraysEqual(
+      self.assertAllClose(
           ids, top_candidate_ids[0, :, :, :num_per_token_logprobs]
       )