Document mask convention

patrick-kidger · Aug 28, 2023 · 6fe42a5 · 6fe42a5
1 parent e771ec8
commit 6fe42a5
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/equinox/nn/_attention.py b/equinox/nn/_attention.py
@@ -236,7 +236,8 @@ def __call__(
             `(kv_seq_length, value_size)`.
         - `mask`: Optional mask preventing attention to certain positions. Should either
             be a JAX array of shape `(query_seq_length, kv_seq_length)`, or (for custom
-            per-head masking) `(num_heads, query_seq_length, kv_seq_length)`.
+            per-head masking) `(num_heads, query_seq_length, kv_seq_length)`. A value of
+            `False` at a position indicates that position should be ignored.
         - `key`: A `jax.random.PRNGKey` used for dropout. Unused if `dropout = 0`.
             (Keyword only argument.)
         - `inference`: As [`equinox.nn.Dropout.__call__`][]. (Keyword only