From 0d387e0839712fa35dbb8bb349de46d736978d5f Mon Sep 17 00:00:00 2001
From: kaixih <kaixih@nvidia.com>
Date: Mon, 15 Jul 2024 17:30:54 +0000
Subject: [PATCH] Update jax doc sdpa

---
 docs/jax.nn.rst | 1 +
 1 file changed, 1 insertion(+)

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