opendatalab/MinerU · error · ValueError

Attention mask batch size {attention_mask.shape[0]} does not

Error message

Attention mask batch size {attention_mask.shape[0]} does not match embeddings batch size {batch_size}

What it means

ValueError raised in _create_bidirectional_mask() when attention_mask.shape[0] differs from the batch size of inputs_embeds (or encoder_hidden_states when provided). The mask is broadcast per sample, so a mismatched leading dimension means the mask does not describe the same batch and expansion would either fail or silently misalign.

Source

Thrown at mineru/model/layout/pp_doclayoutv2.py:154

    inputs_embeds: torch.Tensor,
    attention_mask: Optional[torch.Tensor],
    encoder_hidden_states: Optional[torch.Tensor] = None,
) -> Optional[torch.Tensor]:
    if attention_mask is None:
        return None
    if attention_mask.ndim == 4:
        return attention_mask
    if attention_mask.ndim != 2:
        raise ValueError(
            f"PP-DocLayoutV2 reading-order mask must be 2D or 4D, got shape {tuple(attention_mask.shape)}"
        )

    embeds = encoder_hidden_states if encoder_hidden_states is not None else inputs_embeds
    batch_size, query_length = inputs_embeds.shape[:2]
    key_length = attention_mask.shape[1]

    if attention_mask.shape[0] != batch_size:
        raise ValueError(
            f"Attention mask batch size {attention_mask.shape[0]} does not match embeddings batch size {batch_size}"
        )

    expanded_mask = attention_mask[:, None, None, :].expand(batch_size, 1, query_length, key_length)
    expanded_mask = expanded_mask.to(device=embeds.device, dtype=embeds.dtype)
    min_value = torch.finfo(embeds.dtype).min
    return torch.where(
        expanded_mask > 0,
        torch.zeros(1, dtype=embeds.dtype, device=embeds.device),
        torch.full((1,), min_value, dtype=embeds.dtype, device=embeds.device),
    )


def _load_preprocess_config(model_dir: str) -> Dict:
    config_path = os.path.join(model_dir, "preprocessor_config.json")
    if not os.path.exists(config_path):
        return {}
    with open(config_path, "r", encoding="utf-8") as f:

View on GitHub (pinned to 4fe4bde114)

Solutions

Build attention_mask from the same batch dimension as inputs_embeds: torch.ones(inputs_embeds.shape[0], seq_len, ...).
Slice the mask together with the inputs: mask = mask[:inputs_embeds.shape[0]].
Regenerate the mask per batch inside the inference loop instead of hoisting it out.

Example fix

# before
mask = torch.ones(1, 128)
outputs = model(inputs_embeds=emb_of_batch_8, mask=mask)  # ValueError

# after
mask = torch.ones(emb_of_batch_8.shape[0], 128)
outputs = model(inputs_embeds=emb_of_batch_8, mask=mask)

Defensive patterns

Strategy: validation

Validate before calling

def make_mask(batch_size: int, seq_len: int, device, lengths=None):
    import torch
    mask = torch.zeros(batch_size, seq_len, dtype=torch.long, device=device)
    if lengths is None:
        mask[:] = 1
    else:
        for i, n in enumerate(lengths):
            mask[i, :n] = 1
    return mask

mask = make_mask(inputs_embeds.shape[0], seq_len, inputs_embeds.device)

Type guard

def mask_batch_matches(mask, inputs_embeds) -> bool:
    return mask is not None and mask.shape[0] == inputs_embeds.shape[0]

Prevention

Derive the mask's batch dim from the same tensor you pass as inputs_embeds.
When slicing batches, slice the mask identically.
Never hoist a mask out of the batch loop across batch-size changes.

When it happens

Trigger: inputs_embeds with batch 4 but a 2D mask built for batch 1 (or vice versa); reusing a cached mask after changing batch size; passing a mask computed for encoder_hidden_states of a different batch than inputs_embeds.

Common situations: Batched inference where the mask is generated once for a sample count that later changes; slicing inputs (inputs_embeds[:2]) without slicing the mask; padding logic that grows inputs but not masks.

Related errors

AI-assisted analysis of opendatalab/MinerU@4fe4bde114 (2026-08-14). Data as JSON: /api/errors/3ccbbcf9dc43f2c6. Report an issue: GitHub.