sgl-project/sglang · error · ValueError

pair_postprocess must return the same shape as input

Error message

pair_postprocess must return the same shape as input

What it means

The pair_postprocess hook must preserve the input pair tensor's shape (it modifies values, not structure, of the (num_steps, 2) pairs). Returning a differently-shaped tensor breaks the cached pair schedule invariants.

Source

Thrown at python/sglang/multimodal_gen/runtime/models/schedulers/flow_match_pair.py:521

            sigma_to = self.timestep_to_sigma(timestep_to)
        prev_sample = sample + model_output * (sigma_to - sigma_from)
        return prev_sample

    def _refresh_pair_cache(self) -> None:
        if self.timesteps is None or self.sigmas is None:
            raise RuntimeError("Scheduler not initialized; call set_timesteps() first")

        def _apply_postprocess(pairs: torch.Tensor, source: str) -> torch.Tensor:
            if self._pair_postprocess_fn is None:
                return pairs
            if self._pair_postprocess_requires_source:
                modified = self._pair_postprocess_fn(pairs, source=source)
            else:
                modified = self._pair_postprocess_fn(pairs)
            if not isinstance(modified, torch.Tensor):
                raise TypeError("pair_postprocess must return a torch.Tensor")
            if modified.shape != pairs.shape:
                raise ValueError("pair_postprocess must return the same shape as input")
            return modified

        base_pairs_timesteps = self._make_pairs_from_vector(self.timesteps)
        base_pairs_sigmas = self._make_pairs_from_vector(self.sigmas)

        self.pair_timesteps = _apply_postprocess(base_pairs_timesteps, "timesteps")
        self.pair_sigmas = _apply_postprocess(base_pairs_sigmas, "sigmas")


EntryClass = FlowMatchPairScheduler

View on GitHub (pinned to 0132848349)

Solutions

  1. Return a tensor with identical .shape; if you modify columns, keep torch.stack([col0, col1], dim=1) as _dual_sigma_shift does
  2. Log/assert modified.shape == pairs.shape in your hook during development
  3. Model your hook on the built-in _dual_sigma_shift implementation

Example fix

# before
def pp(pairs, source):
    return pairs[:, 0] * 2  # wrong shape
# after
def pp(pairs, source):
    return torch.stack([pairs[:, 0] * 2, pairs[:, 1] * 2], dim=1)
Defensive patterns

Strategy: validation

Validate before calling

def pp(pairs, source):
    out = modify(pairs)
    assert out.shape == pairs.shape, (out.shape, pairs.shape)
    return out

Prevention

When it happens

Trigger: A postprocess that slices (pairs[:, 0]), stacks along the wrong dim, adds/removes steps, or returns the unstacked columns instead of the stacked pair tensor.

Common situations: Porting a single-vector sigma shift function that operates on 1D sigmas and forgetting to re-stack to (N, 2); off-by-one slicing to drop the terminal sigma.

Related errors


AI-assisted analysis of sgl-project/sglang@0132848349 (2026-08-28). Data as JSON: /api/errors/8e36828566446b42. Report an issue: GitHub.