What causes "Attention mask should be 2D" in dotnet/machinelearning?

Pass the tokenizer/model output attention_mask directly — it is already [batch, seq_len] 2D.

dotnet/machinelearning · error · ArgumentException

Attention mask should be 2D

Error message

Attention mask should be 2D

What it means

AttentionMaskConverter.Create4DCausalAttentionMask throws ArgumentException when the supplied attentionMask tensor is not 2D. The converter expects a [batch, seq_len] padding mask which it expands to [batch, 1, tgt_len, src_len]; higher- or lower-rank inputs cannot be expanded.

Solutions

Pass the tokenizer/model output attention_mask directly — it is already [batch, seq_len] 2D.
If your mask is 4D, use it as-is without re-expanding (skip Create4DCausalAttentionMask).
Reshape 1D masks to 2D with mask.unsqueeze(0) (adding batch dimension) before the call.
Assert mask.ndim == 2 in your forward code before invoking the converter.

Example fix

// before
var mask4d = AttentionMaskConverter.Create4DCausalAttentionMask(mask4dFromPreviousStep, inputShape, dType, kvLen, device);
// after
var mask2d = attentionMaskFromTokenizer; // [batch, seq_len]
var mask4d = AttentionMaskConverter.Create4DCausalAttentionMask(mask2d, inputShape, dType, kvLen, device);

Defensive patterns

Strategy: validation

Validate before calling

if (attentionMask is not null && attentionMask.ndim != 2)
    throw new InvalidOperationException($"attention_mask must be [batch, seq_len] 2D, got ndim={attentionMask.ndim}");

Type guard

bool Is2DMask(Tensor? m) => m is not null && m.ndim == 2;

Try / catch

try { var mask = AttentionMaskConverter.Create4DCausalAttentionMask(attnMask, inputShape, dType, kvLen, device); }
catch (ArgumentException ex) when (ex.Message.Contains("should be 2D")) { logger.LogError("Pass the raw tokenizer 2D attention_mask, not an expanded one"); throw; }

Prevention

Pass the tokenizer output attention_mask directly; never a previously expanded 4D mask.
Unsqueeze 1D masks to add the batch dimension before the call.
Skip Create4DCausalAttentionMask if you already hold a 4D mask.
Assert mask shape in debug builds.

When it happens

Trigger: Passing an already-expanded 4D mask, a 3D per-head mask, or a flattened 1D mask into Create4DCausalAttentionMask.

Common situations: Feeding an output mask from a previous layer back in; reshaping mistakes when batching; porting PyTorch code that keeps masks 4D throughout, while this API wants the raw 2D mask each call.

Understand the failure class

Background: Tensor shape mismatch errors ("must have shape", "expected shape ... got ..."): when tensor dimensions disagree with what an op or layer was told to expect — this error's family across 6 libraries.

Related errors

AI-assisted analysis of dotnet/machinelearning@7b76e69cf9 (2026-09-11). Data as JSON: /api/errors/67ddc776499b378e. Report an issue: GitHub.

Appendix: source

Thrown at src/Microsoft.ML.GenAI.Core/Utility/AttentionMaskConverter.cs:156

    /// <param name="slidingWindow">The sliding window size.</param>
    /// <param name="inputShape">The input shape should be a tuple that defines `(batch_size, query_length)`.</param>
    public static Tensor? Create4DCausalAttentionMask(
        Tensor? attentionMask,
        long[] inputShape,
        ScalarType dType,
        Device device,
        int pastKeyValuesLength = 0,
        int? slidingWindow = null)
    {
        var converter = new AttentionMaskConverter(isCausal: true, slidingWindow: slidingWindow);
        var batchSize = (int)inputShape[0];
        var queryLength = (int)inputShape[1];
        var keyValueLength = pastKeyValuesLength + queryLength;
        if (attentionMask is not null)
        {
            if (attentionMask.ndim != 2)
            {
                throw new ArgumentException("Attention mask should be 2D");
            }
            return converter.To4D(attentionMask, (int)inputShape[1], dType, keyValueLength);
        }

        return converter.ToCausal4D(batchSize, queryLength, keyValueLength, dType, device);
    }

    public static Tensor ExpandMask(
        Tensor mask,
        ScalarType dType,
        int? tgtLen = null)
    {
        var bsz = (int)mask.shape[0];
        var srcLen = (int)mask.shape[1];
        tgtLen ??= srcLen;

        var expandedMask = mask.unsqueeze(1).unsqueeze(1).expand(bsz, 1, tgtLen.Value, srcLen).to(dType);
        var invertedMask = 1.0 - expandedMask;

View on GitHub (pinned to 7b76e69cf9)