diff --git a/docs/jax.nn.rst b/docs/jax.nn.rst
index 33223ee75..adb13f899 100644
--- a/docs/jax.nn.rst
+++ b/docs/jax.nn.rst
@@ -52,3 +52,4 @@ Other functions
     logsumexp
     standardize
     one_hot
+    dot_product_attention