nn.dot_product_attention totally different from manual

[ZagButNoZig]

Hey,

I was playing around with dot product attention and noticed the results are completely different than what I get when manually implementing it. I am probably missing something obvious but can not figure out what:

My code:

x = jax.random.uniform(key=jax.random.key(0), shape=(10, 4))
print(jax.nn.softmax(x @ x.T * 1/math.sqrt(x.shape[-1])) @ x)
print(jax.nn.dot_product_attention(x,x,x))

and I get:

[[0.60897803 0.5603086  0.55994844 0.5347835 ]
 [0.5917675  0.5265384  0.55840343 0.54622614]
 [0.6154469  0.51979136 0.5865181  0.51311886]
 [0.5925973  0.5447331  0.58202934 0.55935323]
 [0.61490166 0.54570115 0.5664482  0.5163673 ]
 [0.5688795  0.56330824 0.56260014 0.56730366]
 [0.59281075 0.5456244  0.58066475 0.55371606]
 [0.57819706 0.5293594  0.5457511  0.55057955]
 [0.5922124  0.5419845  0.5614486  0.5438138 ]
 [0.60580194 0.5260329  0.56715655 0.5314961 ]]

[[0.947667   0.9785799  0.33229148 0.46866846]
 [0.5698887  0.16550303 0.3101946  0.68948054]
 [0.74676657 0.17101455 0.9853538  0.02528262]
 [0.6400418  0.56269085 0.8992138  0.93453753]
 [0.8341402  0.7256162  0.5098531  0.02765214]
 [0.03148878 0.9580188  0.5188192  0.79221416]
 [0.5522419  0.6113529  0.8931755  0.75499094]
 [0.21164179 0.22934973 0.06192303 0.6189376 ]
 [0.49173605 0.54509854 0.41268063 0.5281111 ]
 [0.80325234 0.20560277 0.49416995 0.44019997]]