apache/beam · error · TypeError

batch type and element type must have equivalent dtypes

Error message

batch type and element type must have equivalent dtypes (batch={batch_type.dtype}, element={element_type.dtype})

What it means

Apache Beam's PytorchBatchConverter.from_typehints builds a batch converter from an element type hint and a batch type hint. When batch_type was given explicitly as a PytorchTensor[...] constraint, Beam verifies that its declared dtype matches the element type's dtype; a mismatch means batches and the elements they contain would claim incompatible dtypes, so Beam refuses to construct the converter.

Solutions

1. Make the PytorchTensor batch type dtype identical to the element type dtype, e.g. PytorchTensor[torch.float32, (N, 128)] with element hint PytorchTensor[torch.float32, (128,)].
2. Pass batch_type=torch.Tensor so Beam derives the batch dtype from the element type automatically (PytorchTensor[element_type.dtype, (N,)]).
3. Check for tensor casts in your pipeline (.float(), .double(), .half()) that changed the runtime dtype away from the declared hint.
4. Print both hints (repr) and compare dtype fields before registering the converter.

Example fix

// before
with element_type(PytorchTensor[torch.float32, (128,)]) and \
     batch_type(PytorchTensor[torch.float64, (N, 128)]):
    ...
// after
with element_type(PytorchTensor[torch.float32, (128,)]) and \
     batch_type(PytorchTensor[torch.float32, (N, 128)]):  # dtype matches element
    ...

def check_batch_dtype(batch_type, element_type):
    bt = batch_type if hasattr(batch_type, 'dtype') else PytorchTensor[element_type.dtype, (N,)]
    if bt.dtype != element_type.dtype:
        raise TypeError(f'dtype mismatch: batch={bt.dtype}, element={element_type.dtype}')

def is_dtype_consistent(batch_type, element_type) -> bool:
    return getattr(batch_type, 'dtype', None) == getattr(element_type, 'dtype', None)

When it happens

Trigger: Calling from_typehints (e.g. via @with_batch_type or registering a 'pytorch' BatchConverter) where batch_type=PytorchTensor[torch.float64, (N, 128)] but element_type resolves to a different dtype such as torch.float32 — typically because the PytorchTensor dtype was written explicitly while the element type was inferred from data.

Common situations: Copy-pasting a PytorchTensor hint with a dtype copied from a model checkpoint (float64/float32 mismatch), changing tensor dtype in a DoFn (e.g. .double()) without updating the batch hint, or mixing hints defined in different places of the pipeline.

Related errors

• Batch does not have expected dtype
• Could not align batch type's batch dimension with element…
• 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/cbc20871b584b758. Report an issue: GitHub.

Appendix: source

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

    self.element_shape = element_shape
    self.partition_dimension = partition_dimension

  @staticmethod
  @BatchConverter.register(name="pytorch")
  def from_typehints(element_type,
                     batch_type) -> Optional['PytorchBatchConverter']:
    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)

View on GitHub (pinned to 12126d8942)