MakeAttentionMask

Creates mask for attention weights (Flax utility).

Abstract Signature:

MakeAttentionMask(query_input: Tensor, key_input: Tensor)

PyTorchFlax NNX

PyTorch

API: —

Strategy: Plugin (attention_mask_generator)

Flax NNX

API: flax.nnx.make_attention_mask

Strategy: Direct Mapping

Official Docs ↗