apache/beam · error · TypeError

Batch does not have expected shape

Error message

Batch {batch!r} does not have expected shape: {self.shape!r}

What it means

PytorchTypeConstraint.type_check checks each declared dimension of the PytorchTensor[dtype, shape] hint against the tensor's actual shape. Dimensions equal to the special placeholder N are exempt (variable batch size); every other dimension must match exactly. Beam raises this TypeError when a non-N dimension of the tensor differs from the hinted shape.

Solutions

1. Pad or truncate tensors to a fixed size so every non-N dimension matches the hint (e.g. pad sequences to max_length).
2. Mark the truly variable dimension with N in the hint — any other varying dimension is not allowed by this check.
3. Fix the shape declaration to the real tensor shape (print batch.shape and update the hint accordingly).
4. Reshape with tensor.view(...)/reshape(...) before returning so the tensor matches the expected shape.

Example fix

// before
PytorchTensor[torch.float32, (N, 128)]  # actual tensors vary in seq_len
// after
# pad to a fixed length first
padded = torch.nn.utils.rnn.pad_sequence(tensors, batch_first=True)
# hint: PytorchTensor[torch.float32, (N, 512)] with fixed max_len=512

def check_shape(t: torch.Tensor, shape) -> bool:
    return len(t.shape) == len(shape) and all(
        s == N or t.shape[i] == s for i, s in enumerate(shape))

def matches_hint(t: torch.Tensor, hint) -> bool:
    return isinstance(t, torch.Tensor) and check_shape(t, hint.shape)

try:
    yield batch
except TypeError as e:
    if 'does not have expected shape' in str(e):
        yield pad_to_shape(batch, expected_shape)  # pad/truncate then re-emit
    else:
        raise

When it happens

Trigger: Hinting PytorchTensor[torch.float32, (N, 128)] but emitting tensors whose second dimension is not 128 (e.g. variable-length sequences padded to different lengths), when the variable dimension is not the one marked N, or when extra/missing dimensions shift indices.

Common situations: Variable-length text/audio sequences padded per batch, images with differing resolutions, model input reshaping mistakes, and forgetting that only the N dimension may vary between batches.

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

• Batch does not have expected dtype
• Batch is not an instance of torch.Tensor
• Could not align batch type's batch dimension with element…
• A has been supplied to the model handler, but the required…
• Batch does not have expected shape

AI-assisted analysis of apache/beam@12126d8942 (2026-09-13). Data as JSON: /api/errors/bcd40d4de839b227. Report an issue: GitHub.

Appendix: source

Thrown at sdks/python/apache_beam/typehints/pytorch_type_compatibility.py:106

    return batch.nelement() * batch.element_size()


class PytorchTypeHint():
  class PytorchTypeConstraint(typehints.TypeConstraint):
    def __init__(self, dtype, shape=()):
      self.dtype = dtype
      self.shape = shape

    def type_check(self, batch):
      if not isinstance(batch, torch.Tensor):
        raise TypeError(f"Batch {batch!r} is not an instance of torch.Tensor")
      if not batch.dtype == self.dtype:
        raise TypeError(
            f"Batch {batch!r} does not have expected dtype: {self.dtype!r}")

      for dim in range(len(self.shape)):
        if not self.shape[dim] == N and not batch.shape[dim] == self.shape[dim]:
          raise TypeError(
              f"Batch {batch!r} does not have expected shape: {self.shape!r}")

    def _consistent_with_check_(self, sub):
      # TODO Check sub against batch type, and element type
      return True

    def __key(self):
      return (self.dtype, self.shape)

    def __eq__(self, other) -> bool:
      if isinstance(other, PytorchTypeHint.PytorchTypeConstraint):
        return self.__key() == other.__key()

      return NotImplemented

    def __hash__(self) -> int:
      return hash(self.__key())

View on GitHub (pinned to 12126d8942)