Aligns Softmax masking behavior with JAX for fully masked axis (#21149)

* Fixes softmax masking logic to match JAX behavior

* fix comment

* use backend.numpy.multipy for element-wise multiplication

Author: JyotinderSingh

keras/src/layers/activations/softmax.py

@@ -58,15 +58,25 @@ def call(self, inputs, mask=None):
             inputs += adder
         if isinstance(self.axis, (tuple, list)):
             if len(self.axis) > 1:
-                return backend.numpy.exp(
+                outputs = backend.numpy.exp(
                     inputs
                     - backend.math.logsumexp(
                         inputs, axis=self.axis, keepdims=True
                     )
                 )
             else:
-                return activations.softmax(inputs, axis=self.axis[0])
-        return activations.softmax(inputs, axis=self.axis)
+                outputs = activations.softmax(inputs, axis=self.axis[0])
+        else:
+            outputs = activations.softmax(inputs, axis=self.axis)
+
+        if mask is not None:
+            # Apply the mask to the softmax output to ensure that masked
+            # values are set to 0 in case the entire axis is masked.
+            outputs = backend.numpy.multiply(
+                outputs, backend.cast(mask, outputs.dtype)
+            )
+
+        return outputs
 
     def get_config(self):
         config = super().get_config()

keras/src/layers/activations/softmax_test.py

@@ -49,3 +49,40 @@ def test_softmax_correctness_with_axis(self):
         )
         result = softmax_layer(input)
         self.assertAllClose(result, expected_output)
+
+    def test_softmax_masked_values_are_zero_including_fully_masked(self):
+        """
+        Tests softmax with mask on default axis (-1).
+        Ensures output is 0 where mask is False.
+        Includes a row where all elements are masked.
+        """
+        softmax_layer = softmax.Softmax()  # Default axis = -1
+
+        input = np.array(
+            [
+                [1.0, 2.0, 5.0, 1.0],
+                [1.0, 1.0, 1.0, 1.0],
+                [3.0, 1.0, 2.0, 4.0],
+            ],
+            dtype=np.float32,
+        )
+        mask = np.array(
+            [
+                [True, True, False, False],  # Partially masked
+                [False, False, False, False],  # Fully masked
+                [True, True, True, True],  # Not masked
+            ],
+            dtype=bool,
+        )
+
+        expected_output = np.array(
+            [
+                [0.268941, 0.731059, 0.0, 0.0],  # last two masked
+                [0.0, 0.0, 0.0, 0.0],  # Fully masked row should be all zeros
+                [0.236883, 0.032059, 0.087144, 0.643914],
+            ]
+        )
+
+        result = softmax_layer(input, mask=mask)
+
+        self.assertAllClose(result, expected_output)