Feature Request: Add the is_causal kwarg to flax.nnx.dot_product_attention.

[ibbyml]

Hi Flax community.

I’d like to request adding an is_causal argument to flax.nnx.dot_product_attention for API parity. Right now, flax.nnx.dot_product_attention expects the user to build a causal mask manually (e.g. via the mask utilities in the same module).

Both jax.nn.dot_product_attention and PyTorch’s torch.nn.functional.scaled_dot_product_attention support the is_causal argument.

Rough Proposal

Extend the NNX function signature with a keyword-only argument:

def dot_product_attention(..., is_causal: bool = False):
    ...

Behavior wise, you'd have 2 paths:

Jax.nn Path: if dropout is not activated and module=None, flax.nnx.dot_product_attention will already call jax.nn.dot_product_attention under the hood. We can simply forward the argument:

jax.nn.dot_product_attention(query, key, value, bias, mask, is_causal=is_causal)

Manual Path: When the manual path is used (dropout active or module non-None), we could interpret is_causal=True as “apply a standard causal mask on top of any existing mask”, using the existing mask utilities to construct and combine it.
Conceptually:

combined_mask = mask
if is_causal:
    causal_mask = ...  # standard [batch..., 1, q_length, kv_length] mask
    combined_mask = combine_masks(combined_mask, causal_mask)

If you’d prefer this to be limited to the jax.nn path only (or not added at all), I’m happy to adjust or drop the idea.

If this sounds reasonable, I can open a small PR with the minimal implementation and tests.

Thanks for considering it!

[vfdev-5]

Thanks for suggestion @ibbyml ! I think it makes sense to add is_causal argument and support it in both paths.

[ibbyml]

Great. I'll get started on opening the pull request.