From 5f42e4b501296733eee55c3ad03bdeb7d57ce6ac Mon Sep 17 00:00:00 2001 From: V-MoE Authors Date: Wed, 30 Sep 2026 21:42:31 -0700 Subject: [PATCH] Fix non-deterministic ordering assumption in input_pipeline_test. PiperOrigin-RevId: 991418554 --- vmoe/data/input_pipeline_test.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/vmoe/data/input_pipeline_test.py b/vmoe/data/input_pipeline_test.py index 005abb5..ac73d16 100644 --- a/vmoe/data/input_pipeline_test.py +++ b/vmoe/data/input_pipeline_test.py @@ -119,9 +119,10 @@ def test_eval(self): # Eval data has a '__valid__' field, since fake examples were added. self.assertSetEqual(set(data[0].keys()), {'x', 'y', input_pipeline.VALID_KEY}) - # Eval data is not shuffled, and the fake data corresponds to the first - # element. - self.assertTupleEqual(tuple(data[0]['x']), (0, 1, 2, 0, 0)) + # Eval data has valid elements (order may vary due to non-deterministic + # parallel mapping) and fake elements at the end. + self.assertSequenceEqual(sorted(data[0]['x'][:3]), (0, 1, 2)) + self.assertTupleEqual(tuple(data[0]['x'][3:]), (0, 0)) self.assertTupleEqual(tuple(data[0][input_pipeline.VALID_KEY]), (True, True, True, False, False))