diff --git a/site/en/guide/extension_type.ipynb b/site/en/guide/extension_type.ipynb
index 76e20e8d283..7e8edeea7c9 100644
--- a/site/en/guide/extension_type.ipynb
+++ b/site/en/guide/extension_type.ipynb
@@ -1822,13 +1822,17 @@
         "                  transpose_a=False, transpose_b=False,\n",
         "                  adjoint_a=False, adjoint_b=False,\n",
         "                  a_is_sparse=False, b_is_sparse=False,\n",
-        "                  output_type=None):\n",
+        "                  output_type=None,\n",
+        "                  grad_a=False, grad_b=False,\n",
+        "                  name=None,\n",
+        "                  ):\n",
         "  if isinstance(a, MaskedTensor):\n",
         "    a = a.with_default(0)\n",
         "  if isinstance(b, MaskedTensor):\n",
         "    b = b.with_default(0)\n",
         "  return tf.matmul(a, b, transpose_a, transpose_b, adjoint_a,\n",
-        "                   adjoint_b, a_is_sparse, b_is_sparse, output_type)"
+        "                   adjoint_b, a_is_sparse, b_is_sparse,\n",
+        "                   output_type)"
       ]
     },
     {