Update softmax check tests

Author: Frightera

keras/engine/training_test.py

@@ -18,7 +18,6 @@
 import io
 import sys
 import tempfile
-import warnings
 
 import numpy as np
 import tensorflow.compat.v2 as tf
@@ -5010,31 +5009,23 @@ def test_sequential_model_output(self):
                     layers_module.Dense(1, activation=activation),
                 ]
             )
-            with warnings.catch_warnings(record=True) as w:
-                warnings.simplefilter("always")
+            with self.assertRaisesRegex(
+                ValueError,
+                "has a single unit output, but the activation is softmax.*",
+            ):
                 model.compile()
-                self.assertIs(w[-1].category, SyntaxWarning)
-                self.assertIn(
-                    "Found a layer with softmax activation and single unit "
-                    "output",
-                    str(w[-1].message),
-                )
             del model
 
     def test_functional_model_output(self):
         inputs = input_layer.Input(shape=(10,))
         for activation in ["softmax", tf.nn.softmax, layers_module.Softmax()]:
             x = layers_module.Dense(1, activation=activation)(inputs)
             model = training_module.Model(inputs, x)
-            with warnings.catch_warnings(record=True) as w:
-                warnings.simplefilter("always")
+            with self.assertRaisesRegex(
+                ValueError,
+                "has a single unit output, but the activation is softmax.*",
+            ):
                 model.compile()
-                self.assertIs(w[-1].category, SyntaxWarning)
-                self.assertIn(
-                    "Found a layer with softmax activation and single unit "
-                    "output",
-                    str(w[-1].message),
-                )
             del model
 
     def test_multi_output_model(self):
@@ -5043,15 +5034,11 @@ def test_multi_output_model(self):
             x = layers_module.Dense(1, activation=activation)(inputs)
             y = layers_module.Dense(1, activation=activation)(inputs)
             model = training_module.Model(inputs, [x, y])
-            with warnings.catch_warnings(record=True) as w:
-                warnings.simplefilter("always")
+            with self.assertRaisesRegex(
+                ValueError,
+                "has a single unit output, but the activation is softmax.*",
+            ):
                 model.compile()
-                self.assertIs(w[-1].category, SyntaxWarning)
-                self.assertIn(
-                    "Found a layer with softmax activation and single unit "
-                    "output",
-                    str(w[-1].message),
-                )
             del model
 
     def test_multi_input_output_model(self):
@@ -5063,15 +5050,11 @@ def test_multi_input_output_model(self):
             x = layers_module.Dense(1, activation=activation)(inputs[0])
             y = layers_module.Dense(1, activation=activation)(inputs[1])
             model = training_module.Model(inputs, [x, y])
-            with warnings.catch_warnings(record=True) as w:
-                warnings.simplefilter("always")
+            with self.assertRaisesRegex(
+                ValueError,
+                "has a single unit output, but the activation is softmax.*",
+            ):
                 model.compile()
-                self.assertIs(w[-1].category, SyntaxWarning)
-                self.assertIn(
-                    "Found a layer with softmax activation and single unit "
-                    "output",
-                    str(w[-1].message),
-                )
             del model