apache/beam · error · TypeError

Failed to align batch type's batch dimension with element ty

Error message

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

What it means

The batch type's shape must contain exactly one N (the batch/partition dimension) and removing it must yield the element type's shape. If the remaining dimensions don't equal element_type.shape, shapes cannot be aligned.

Source

Thrown at sdks/python/apache_beam/typehints/batch.py:210

        raise TypeError("Element type is not a dtype") from e

    if not isinstance(batch_type, NumpyTypeHint.NumpyTypeConstraint):
      if not batch_type == np.ndarray:
        raise TypeError(
            "batch type must be np.ndarray or "
            "beam.typehints.batch.NumpyArray[..]")
      batch_type = NumpyArray[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(
          "Failed to align batch type's batch dimension with element type. "
          f"(batch type dimensions: {batch_type.shape}, element type "
          f"dimenstions: {element_type.shape}")

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

  def produce_batch(self, elements):
    return np.stack(elements, axis=self.partition_dimension)

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

View on GitHub (pinned to 12126d8942)

Solutions

1. Ensure batch shape is element shape with a single N inserted at the batching axis, e.g. element (3,) -> batch (N, 3)
2. Check the shape you pass to NumpyArray[dtype, shape] counts all element dims
3. Use () element shape for 1-D batches of scalars

Example fix

// before
BatchConverter.from_typehints(NumpyArray[np.float64, (3,)], NumpyArray[np.float64, (N,)])
// after
BatchConverter.from_typehints(NumpyArray[np.float64, (3,)], NumpyArray[np.float64, (N, 3)])

# element shape (3,) -> batch shape (N, 3)
assert batch_shape.count(N) == 1
computed = tuple(d for d in batch_shape if d != N)
assert computed == element_shape

def shapes_align(batch_shape, elem_shape, N):
    if batch_shape.count(N) != 1: return False
    dims = list(batch_shape); dims.remove(N)
    return tuple(dims) == tuple(elem_shape)

When it happens

Trigger: from_typehints(NumpyArray[np.float64, (3,)], NumpyArray[np.float64, (N,)]) — batch is 1-D but element is 1-D with size 3, removing N gives () which != (3,); or batch hint shape with wrong ordering of dims around N.

Common situations: Declaring scalar elements but batching into 2-D arrays or vice versa; hand-writing NumpyArray shapes where the N placement doesn't match the element shape.

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 {batch!r} does not have expected shape: {self.shape!r}
• Element type is not a dtype
• batch type must be np.ndarray or beam.typehints.batch.NumpyA
• batch type and element type must have equivalent dtypes (bat
• Batch {batch!r} is not an instance of ndarray

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