apache/beam · error · TypeError

Could not align batch type's batch dimension with element…

Error message

Could not align batch type's batch dimension with element type. (batch type dimensions: {batch_type.shape}, element type dimenstions: {element_type.shape}

What it means

from_typehints removes the N (batch dimension) placeholder from the batch type's shape and requires the remaining dimensions to exactly equal the element type's shape. If dropping the partition dimension from batch_type.shape does not reproduce element_type.shape, the batch and element types do not describe the same underlying tensors, so the converter cannot be built.

Solutions

1. Ensure batch_type.shape is exactly element_type.shape with a single N inserted at the partition dimension, e.g. element (128,) -> batch (N, 128).
2. If elements are scalars, use batch_type=PytorchTensor[dtype, (N,)] and element_type=PytorchTensor[dtype, ()].
3. Remove redundant/extra dimensions from the batch hint so after dropping N the tuples match exactly.
4. Pass plain torch.Tensor as batch_type so Beam constructs PytorchTensor[element.dtype, (N,)] consistent with the element type.

Example fix

// before
batch_type=PytorchTensor[torch.float32, (N, 224, 224, 3)],
element_type=PytorchTensor[torch.float32, (3, 224, 224)]
// after
batch_type=PytorchTensor[torch.float32, (N, 3, 224, 224)],
element_type=PytorchTensor[torch.float32, (3, 224, 224)]

def check_batch_shape(batch_type, element_type):
    dims = [d for d in batch_type.shape if d != N]
    if tuple(dims) != tuple(element_type.shape):
        raise TypeError(f'shape mismatch: batch={batch_type.shape}, element={element_type.shape}')

def shapes_align(batch_type, element_type) -> bool:
    dims = [d for d in getattr(batch_type, 'shape', ()) if d != N]
    return tuple(dims) == tuple(getattr(element_type, 'shape', ()))

When it happens

Trigger: Registering a pytorch BatchConverter where e.g. batch_type=PytorchTensor[torch.float32, (N, 128)] but element_type=PytorchTensor[torch.float32, (64,)], or the batch shape contains no N at all, or extra/missing dimensions shift the alignment.

Common situations: Typos in dimension sizes after refactoring a model's input shape, forgetting that PytorchTensor with only a dtype defaults the shape to (N,), channels-first vs channels-last ordering differences, or a (N,) batch shape while elements are scalars with shape ().

Related errors

• Batch does not have expected shape
• batch type and element type must have equivalent dtypes
• A has been supplied to the model handler, but the required…
• An Option type-hint only accepts a single type parameter.
• bad type

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

Appendix: source

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

    if not isinstance(element_type, PytorchTypeHint.PytorchTypeConstraint):
      element_type = PytorchTensor[element_type, ()]

    if not isinstance(batch_type, PytorchTypeHint.PytorchTypeConstraint):
      if not batch_type == torch.Tensor:
        raise TypeError(
            "batch type must be torch.Tensor or "
            "beam.typehints.pytorch_type_compatibility.PytorchTensor[..]")
      batch_type = PytorchTensor[element_type.dtype, (N, )]

    if not batch_type.dtype == element_type.dtype:
      raise TypeError(
          "batch type and element type must have equivalent dtypes "
          f"(batch={batch_type.dtype}, element={element_type.dtype})")
    computed_element_shape = list(batch_type.shape)
    partition_dimension = computed_element_shape.index(N)
    computed_element_shape.pop(partition_dimension)
    if not tuple(computed_element_shape) == element_type.shape:
      raise TypeError(
          "Could not align batch type's batch dimension with element type. "
          f"(batch type dimensions: {batch_type.shape}, element type "
          f"dimenstions: {element_type.shape}")

    return PytorchBatchConverter(
        batch_type,
        element_type,
        batch_type.dtype,
        element_type.shape,
        partition_dimension)

  def produce_batch(self, elements):
    return torch.stack(elements, dim=self.partition_dimension)

  def explode_batch(self, batch):
    """Convert an instance of B to Generator[E]."""
    yield from torch.swapaxes(batch, self.partition_dimension, 0)

View on GitHub (pinned to 12126d8942)