{"record":{"id":"75a1818071ddf481","repo":"sgl-project/sglang","slug":"z-image-caption-tensor-must-have-rank-2-or-3","errorCode":null,"errorMessage":"Z-Image caption tensor must have rank 2 or 3","messagePattern":"Z-Image caption tensor must have rank 2 or 3","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"python/sglang/multimodal_gen/runtime/breakable_cuda_graph/model_padders/zimage.py","lineNumber":41,"sourceCode":"\n\ndef _first_caption_tensor(encoder_hidden_states: Any) -> torch.Tensor | None:\n    tensor = bcg_utils.first_tensor(encoder_hidden_states)\n    if not torch.is_tensor(tensor):\n        return None\n    if tensor.dim() == 2:\n        return tensor\n    if tensor.dim() == 3:\n        return tensor[0]\n    return None\n\n\ndef _caption_seq_len(tensor: torch.Tensor) -> int:\n    if tensor.dim() == 2:\n        return int(tensor.shape[0])\n    if tensor.dim() == 3:\n        return int(tensor.shape[1])\n    raise ValueError(\"Z-Image caption tensor must have rank 2 or 3\")\n\n\ndef _pad_caption(obj: Any, *, target: int) -> Any:\n    if torch.is_tensor(obj):\n        if obj.dim() == 2:\n            return bcg_utils.pad_tensor_dim(obj, 0, target)\n        if obj.dim() == 3:\n            return bcg_utils.pad_tensor_dim(obj, 1, target)\n        return obj\n    if isinstance(obj, list):\n        return [_pad_caption(item, target=target) for item in obj]\n    if isinstance(obj, tuple):\n        return tuple(_pad_caption(item, target=target) for item in obj)\n    return obj\n\n\ndef _unwrap_model(current_model: Any) -> Any:\n    for attr in (\"module\", \"_orig_mod\"):","sourceCodeStart":23,"sourceCodeEnd":59,"githubUrl":"https://github.com/sgl-project/sglang/blob/0132848349585cfe6aae51c4941cbae872505f8a/python/sglang/multimodal_gen/runtime/breakable_cuda_graph/model_padders/zimage.py#L23-L59","documentation":"The Z-Image model padder must read the caption sequence length from a prompt tensor: rank 2 ([seq, dim]) yields shape[0], rank 3 ([batch, seq, dim]) yields shape[1]. Any other rank cannot be unambiguously padded, so _caption_seq_len (used by pad_zimage_prompt_kwargs) raises.","triggerScenarios":"Calling pad_zimage_prompt_kwargs with a Z-Image caption tensor whose .dim() is neither 2 nor 3 — e.g. a flat 1D token-id tensor or a 4D input.","commonSituations":"Upstream prompt-encoding changes passing raw token ids instead of embedded 2D/3D captions; a new multimodal input key leaking into the caption padding path; batch-handling changes in the multimodal runtime altering tensor rank.","solutions":["Ensure caption embeddings are passed as [seq, hidden] or [batch, seq, hidden] before padding.","If you have 1D token ids, embed/expand them to rank 2 first (run the text encoder or add a batch dim).","If the tensor should not be treated as a caption, exclude it from the kwargs the Z-Image padder iterates."],"exampleFix":"// before\ncaption = input_ids  # rank 1: [seq]\nkwargs = pad_zimage_prompt_kwargs({\"caption\": caption}, target=128)\n// after\ncaption = text_encoder(input_ids)  # rank 2: [seq, hidden]\nkwargs = pad_zimage_prompt_kwargs({\"caption\": caption}, target=128)","handlingStrategy":"type-guard","validationCode":"def check_caption(t: torch.Tensor) -> None:\n    assert t.dim() in (2, 3), f\"caption must be rank 2/3, got rank {t.dim()}\"\n\ncheck_caption(caption)\nout = pad_zimage_prompt_kwargs(prompt_kwargs, target=target)","typeGuard":"def is_valid_caption(tensor: object) -> bool:\n    return isinstance(tensor, torch.Tensor) and tensor.dim() in (2, 3)","tryCatchPattern":"try:\n    padded = pad_zimage_prompt_kwargs(kwargs, target=t)\nexcept ValueError as e:\n    if \"rank 2 or 3\" in str(e):\n        kwargs[\"caption\"] = embed_to_2d(kwargs[\"caption\"])\n        padded = pad_zimage_prompt_kwargs(kwargs, target=t)\n    else:\n        raise","preventionTips":["Standardize caption tensors to [batch, seq, hidden] before the padding stage.","Log tensor shapes at prompt-build time to catch rank drift early.","Exclude non-caption inputs from kwargs passed to the Z-Image padder."],"tags":["tensor-shape","cuda-graph","z-image","multimodal","padding"],"backgroundTag":"invalid-tensor-rank","analyzedSha":"0132848349585cfe6aae51c4941cbae872505f8a","analyzedAt":"2026-08-28T05:10:05.995Z","schemaVersion":2},"datasetVersion":"2026-08-28T06:17:29.519Z"}