fix test not to depend on details of JAX shape error type/message

PiperOrigin-RevId: 583324422
google-deepmind · Nov 17, 2023 · af9497f · af9497f
1 parent ce05309
commit af9497f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/haiku/_src/nets/vqvae_test.py b/haiku/_src/nets/vqvae_test.py
@@ -91,7 +91,7 @@ def testConstruct(self, constructor, kwargs):
   def testShapeChecking(self, constructor, kwargs):
     vqvae_module = constructor(**kwargs)
     wrong_shape_input = np.random.randn(100, kwargs['embedding_dim'] * 2)
-    with self.assertRaisesRegex(TypeError, 'total size must be unchanged'):
+    with self.assertRaises((ValueError, TypeError)):
       vqvae_module(
           jnp.array(wrong_shape_input.astype(np.float32)), is_training=False)