diff --git a/guides/ipynb/understanding_masking_and_padding.ipynb b/guides/ipynb/understanding_masking_and_padding.ipynb
index 83283c9b39..50da78fdda 100644
--- a/guides/ipynb/understanding_masking_and_padding.ipynb
+++ b/guides/ipynb/understanding_masking_and_padding.ipynb
@@ -526,7 +526,7 @@
     "        broadcast_float_mask = tf.expand_dims(tf.cast(mask, \"float32\"), -1)\n",
     "        inputs_exp = tf.exp(inputs) * broadcast_float_mask\n",
     "        inputs_sum = tf.reduce_sum(\n",
-    "            inputs_exp * broadcast_float_mask, axis=-1, keepdims=True\n",
+    "            inputs_exp * broadcast_float_mask, axis=1, keepdims=True\n",
     "        )\n",
     "        return inputs_exp / inputs_sum\n",
     "\n",
@@ -594,4 +594,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 0
-}
\ No newline at end of file
+}
diff --git a/guides/md/understanding_masking_and_padding.md b/guides/md/understanding_masking_and_padding.md
index db54ddfcb8..47875bb2af 100644
--- a/guides/md/understanding_masking_and_padding.md
+++ b/guides/md/understanding_masking_and_padding.md
@@ -418,7 +418,7 @@ class TemporalSoftmax(keras.layers.Layer):
         broadcast_float_mask = tf.expand_dims(tf.cast(mask, "float32"), -1)
         inputs_exp = tf.exp(inputs) * broadcast_float_mask
         inputs_sum = tf.reduce_sum(
-            inputs_exp * broadcast_float_mask, axis=-1, keepdims=True
+            inputs_exp * broadcast_float_mask, axis=1, keepdims=True
         )
         return inputs_exp / inputs_sum