Repository: google/flax

Issue:
Feature request: optionally sow attention weight in `dot_product_attention`
### Discussed in https://github.com/google/flax/discussions/2869

<div type='discussions-op-text'>

<sup>Originally posted by **JyChang012** February 13, 2023</sup>
It seems that there is no good way to retrieve attention weight using `MultiHeadDotProductAttention`. Probably we can add an optional `mdl` parameter to  `dot_product_attention` and `dot_product_attention_weights`, and `sow` the weight inside. This should have no effect when `intermediates` is not mutable. Here is my [implementation](https://github.com/JyChang012/espnet/blob/jax/asr_train/espnex/models/transformer/attention.py). I can create a PR if you like.</div>

Also requested in #3496

Discussion:
chiamp: cc: @cgarciae @JyChang012 @Xiaoming-Zhao
JyChang012: 👍
cgarciae: Just a small comment, can we give a more descriptive name to the `mdl` argument?
chiamp: > Just a small comment, can we give a more descriptive name to the `mdl` argument?

would `module` work instead
codecov-commenter: ## [Codecov](https://app.codecov.io/gh/google/flax/pull/3529?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=google) Report
Attention: `2 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`512a6d8`)](https://app.codecov.io/gh/google/flax/commit/512a6d883cb59fbf2e80a422f5ed482ba8ab5dba?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=google) 56.16% compared to head [(`ad996e6`)](https://app.codecov.io/gh/google/flax/pull/3529?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=google) 56.05%.
> Report is 2 commits behind head on main.

| [Files](https://app.codecov.io/gh/google/flax/pull/3529?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=google) | Patch % | Lines |
|---|---|---|
| [flax/linen/attention.py](https://app.codecov.io/gh/google/flax/pull/3529?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=google#diff-ZmxheC9saW5lbi9hdHRlbnRpb24ucHk=) | 0.00% | [2 Missing :warning: ](https://app.codecov.io/gh/google/flax/pull/3529?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=google) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main    #3529      +/-   ##
==========================================
- Coverage   56.16%   56.05%   -0.12%     
==========================================
  Files         100      100              
  Lines       11861    11863       +2     
==========================================
- Hits         6662     6650      -12     
- Misses       5199     5213      +14     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/google/flax/pull/3529?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=google).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=google).


Patch:
diff --git a/flax/linen/attention.py b/flax/linen/attention.py
index 99e56552..23f643b4 100644
--- a/flax/linen/attention.py
+++ b/flax/linen/attention.py
@@ -50,6 +50,7 @@ def dot_product_attention_weights(
   deterministic: bool = False,
   dtype: Optional[Dtype] = None,
   precision: PrecisionLike = None,
+  module: Optional[Module] = None,
 ):
   """Computes dot-product attention weights given query and key.
 
@@ -76,6 +77,10 @@ def dot_product_attention_weights(
     dtype: the dtype of the computation (default: infer from inputs and params)
     precision: numerical precision of the computation see `jax.lax.Precision`
       for details.
+    module: the Module that will sow the attention weights into the
+      'intermediates' collection. Remember to mark 'intermediates' as mutable via
+      `mutable=['intermediates'] in order to have that collection returned.
+      If `module` is None, the attention weights will not be sowed.
 
   Returns:
     Output of shape `[batch..., num_heads, q_length, kv_length]`.
@@ -107,6 +112,9 @@ def dot_product_attention_weights(
   # normalize the attention weights
   attn_weights = jax.nn.softmax(attn_weights).astype(dtype)
 
+  if module:
+    module.sow('intermediates', 'attention_weights', attn_weights)
+
   # apply attention dropout
   if not deterministic and dropout_rate > 0.0:
     keep_prob = 1.0 - dropout_rate
@@ -134,6 +142,7 @@ def dot_product_attention(
   deterministic: bool = False,
   dtype: Optional[Dtype] = None,
   precision: PrecisionLike = None,
+  module: Optional[Module] = None,
 ):
   """Computes dot-product attention given query, key, and value.
 
@@ -164,6 +173,10 @@ def dot_product_attention(
     dtype: the dtype of the computation (default: infer from inputs)
     precision: numerical precision of the computation see `jax.lax.Precision`
       for details.
+    module: the Module that will sow the attention weights into the
+      'intermediates' collection. Remember to mark 'intermediates' as mutable via
+      `mutable=['intermediates'] in order to have that collection returned.
+      If `module` is None, the attention weights will not be sowed.
 
   Returns:
     Output of shape `[batch..., q_length, num_heads, v_depth_per_head]`.
@@ -191,6 +204,7 @@ def dot_product_attention(
     deterministic,
     dtype,
     precision,
+    module,
   )
 
   # return weighted sum over values for each query position
@@ -306,6 +320,7 @@ class MultiHeadDotProductAttention(Module):
     mask: Optional[Array] = None,
     deterministic: Optional[bool] = None,
     dropout_rng: Optional[PRNGKey] = None,
+    return_weights: bool = False,
   ):
     ...
 
@@ -318,6 +333,7 @@ class MultiHeadDotProductAttention(Module):
     mask: Optional[Array] = None,
     deterministic: Optional[bool] = None,
     dropout_rng: Optional[PRNGKey] = None,
+    return_weights: bool = False,
   ):
     ...
 
@@ -332,6 +348,7 @@ class MultiHeadDotProductAttention(Module):
     mask: Optional[Array] = None,
     deterministic: Optional[bool] = None,
     dropout_rng: Optional[PRNGKey] = None,
+    return_weights: bool = False,
   ):
     """Applies multi-head dot product attention on the input data.
 
@@ -358,6 +375,10 @@ class MultiHeadDotProductAttention(Module):
         dropout, whereas if true, the attention weights are deterministic.
       dropout_rng: optional rng key to pass to the attention layer's dropout
         mask. Otherwise, self.make_rng('dropout') is used instead.
+      return_weights: if `True`, the attention weights are sowed into the
+        'intermediates' collection. Remember to mark 'intermediates' as
+        mutable via `mutable=['intermediates'] in order to have that
+        collection returned.
 
     Returns:
       output of shape `[batch_sizes..., length, features]`.
@@ -506,18 +527,33 @@ class MultiHeadDotProductAttention(Module):
       m_deterministic = True
 
     # apply attention
-    x = self.attention_fn(
-      query,
-      key,
-      value,
-      mask=mask,
-      dropout_rng=dropout_rng,
-      dropout_rate=self.dropout_rate,
-      broadcast_dropout=self.broadcast_dropout,
-      deterministic=m_deterministic,
-      dtype=self.dtype,
-      precision=self.precision,
-    )  # pytype: disable=wrong-keyword-args
+    if return_weights:
+      x = self.attention_fn(
+        query,
+        key,
+        value,
+        mask=mask,
+        dropout_rng=dropout_rng,
+        dropout_rate=self.dropout_rate,
+        broadcast_dropout=self.broadcast_dropout,
+        deterministic=m_deterministic,
+        dtype=self.dtype,
+        precision=self.precision,
+        module=self if return_weights else None,
+      )  # pytype: disable=wrong-keyword-args
+    else:
+      x = self.attention_fn(
+        query,
+        key,
+        value,
+        mask=mask,
+        dropout_rng=dropout_rng,
+        dropout_rate=self.dropout_rate,
+        broadcast_dropout=self.broadcast_dropout,
+        deterministic=m_deterministic,
+        dtype=self.dtype,
+        precision=self.precision,
+      )
     # back to the original inputs dimensions
     out = DenseGeneral(
       features=features,


Test patch:
diff --git a/flax/experimental/nnx/tests/nn/test_conv.py b/flax/experimental/nnx/tests/nn/test_conv.py
index 2992f7d9..76be459e 100644
--- a/flax/experimental/nnx/tests/nn/test_conv.py
+++ b/flax/experimental/nnx/tests/nn/test_conv.py
@@ -24,18 +24,18 @@ from flax import linen
 from flax.experimental import nnx
 
 
-class TestConvLinenConsistency(parameterized.TestCase):
+class TestLinenConsistency(parameterized.TestCase):
 
   @parameterized.product(
-      strides = [None, (2, 3)],
-      padding = ['VALID', (4, 2)],
-      input_dilation = [(2, 3)],
-      kernel_dilation = [(2, 3)],
-      feature_group_count = [3],
+      strides = [None, 1, (2, 3)],
+      padding = ['SAME', 'VALID', 'CIRCULAR', (4, 2)],
+      input_dilation = [1, (2, 3)],
+      kernel_dilation = [1, (2, 3)],
+      feature_group_count = [1, 3],
       use_bias = [True, False],
-      dtype = [jnp.float32],
-      param_dtype = [jnp.float16],
-      precision = [Precision.HIGHEST],
+      dtype = [jnp.float32, jnp.float16],
+      param_dtype = [jnp.float32, jnp.float16],
+      precision = [Precision.DEFAULT, Precision.HIGH, Precision.HIGHEST],
   )
   def test_nnx_linen_equivalence(self, **kwargs):
     key = jax.random.key(42)
diff --git a/tests/linen/linen_attention_test.py b/tests/linen/linen_attention_test.py
index 44d339da..932f0ec8 100644
--- a/tests/linen/linen_attention_test.py
+++ b/tests/linen/linen_attention_test.py
@@ -343,6 +343,57 @@ class AttentionTest(parameterized.TestCase):
       with self.assertRaises(errors.ScopeParamShapeError):
         module.apply(initial_vars, query, key, causal_mask)
 
+  def test_multihead_sow_attention_weights(self):
+    rng = random.key(0)
+    x = jnp.ones((4, 6, 5))
+
+    class Model(nn.Module):
+      attention_kwargs: dict
+
+      @nn.compact
+      def __call__(self, x, return_weights=False):
+        x = nn.MultiHeadDotProductAttention(**self.attention_kwargs)(
+          x, return_weights=return_weights
+        )
+        x = nn.MultiHeadDotProductAttention(**self.attention_kwargs)(x)
+        x = nn.MultiHeadDotProductAttention(**self.attention_kwargs)(
+          x, return_weights=return_weights
+        )
+        return x
+
+    module = Model(
+      dict(
+        num_heads=8,
+        qkv_features=16,
+        kernel_init=initializers.ones,
+        bias_init=initializers.zeros,
+        deterministic=False,
+      )
+    )
+    v = module.init(rng, x)
+    _, intermediates = module.apply(
+      v, x, mutable=['intermediates'], return_weights=True
+    )
+    self.assertEqual(
+      intermediates['intermediates']['MultiHeadDotProductAttention_0'][
+        'attention_weights'
+      ][0].shape,
+      (4, 8, 6, 6),
+    )
+    self.assertNotIn(
+      'MultiHeadDotProductAttention_1', intermediates['intermediates']
+    )
+    self.assertEqual(
+      intermediates['intermediates']['MultiHeadDotProductAttention_2'][
+        'attention_weights'
+      ][0].shape,
+      (4, 8, 6, 6),
+    )
+    _, intermediates = module.apply(
+      v, x, mutable=['intermediates'], return_weights=False
+    )
+    self.assertNotIn('intermediates', intermediates)
+
 
 if __name__ == '__main__':
   absltest.main()