{"record":{"id":"3ccbbcf9dc43f2c6","repo":"opendatalab/MinerU","slug":"attention-mask-batch-size-attention-mask-shape-0","errorCode":null,"errorMessage":"Attention mask batch size {attention_mask.shape[0]} does not match embeddings batch size {batch_size}","messagePattern":"Attention mask batch size (.+?) does not match embeddings batch size (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"mineru/model/layout/pp_doclayoutv2.py","lineNumber":154,"sourceCode":"    inputs_embeds: torch.Tensor,\n    attention_mask: Optional[torch.Tensor],\n    encoder_hidden_states: Optional[torch.Tensor] = None,\n) -> Optional[torch.Tensor]:\n    if attention_mask is None:\n        return None\n    if attention_mask.ndim == 4:\n        return attention_mask\n    if attention_mask.ndim != 2:\n        raise ValueError(\n            f\"PP-DocLayoutV2 reading-order mask must be 2D or 4D, got shape {tuple(attention_mask.shape)}\"\n        )\n\n    embeds = encoder_hidden_states if encoder_hidden_states is not None else inputs_embeds\n    batch_size, query_length = inputs_embeds.shape[:2]\n    key_length = attention_mask.shape[1]\n\n    if attention_mask.shape[0] != batch_size:\n        raise ValueError(\n            f\"Attention mask batch size {attention_mask.shape[0]} does not match embeddings batch size {batch_size}\"\n        )\n\n    expanded_mask = attention_mask[:, None, None, :].expand(batch_size, 1, query_length, key_length)\n    expanded_mask = expanded_mask.to(device=embeds.device, dtype=embeds.dtype)\n    min_value = torch.finfo(embeds.dtype).min\n    return torch.where(\n        expanded_mask > 0,\n        torch.zeros(1, dtype=embeds.dtype, device=embeds.device),\n        torch.full((1,), min_value, dtype=embeds.dtype, device=embeds.device),\n    )\n\n\ndef _load_preprocess_config(model_dir: str) -> Dict:\n    config_path = os.path.join(model_dir, \"preprocessor_config.json\")\n    if not os.path.exists(config_path):\n        return {}\n    with open(config_path, \"r\", encoding=\"utf-8\") as f:","sourceCodeStart":136,"sourceCodeEnd":172,"githubUrl":"https://github.com/opendatalab/MinerU/blob/4fe4bde114a23ee5dd637eae99b767f4669bf58c/mineru/model/layout/pp_doclayoutv2.py#L136-L172","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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."],"exampleFix":"# before\nmask = torch.ones(1, 128)\noutputs = model(inputs_embeds=emb_of_batch_8, mask=mask)  # ValueError\n\n# after\nmask = torch.ones(emb_of_batch_8.shape[0], 128)\noutputs = model(inputs_embeds=emb_of_batch_8, mask=mask)","handlingStrategy":"validation","validationCode":"def make_mask(batch_size: int, seq_len: int, device, lengths=None):\n    import torch\n    mask = torch.zeros(batch_size, seq_len, dtype=torch.long, device=device)\n    if lengths is None:\n        mask[:] = 1\n    else:\n        for i, n in enumerate(lengths):\n            mask[i, :n] = 1\n    return mask\n\nmask = make_mask(inputs_embeds.shape[0], seq_len, inputs_embeds.device)","typeGuard":"def mask_batch_matches(mask, inputs_embeds) -> bool:\n    return mask is not None and mask.shape[0] == inputs_embeds.shape[0]","tryCatchPattern":null,"preventionTips":["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."],"tags":["pytorch","deep-learning","shape-mismatch","batching","attention-mask"],"backgroundTag":null,"analyzedSha":"4fe4bde114a23ee5dd637eae99b767f4669bf58c","analyzedAt":"2026-08-14T21:29:18.456Z","schemaVersion":2},"datasetVersion":"2026-08-15T22:17:37.221Z"}