apache/beam · error · TypeError

Batch does not have expected dtype

Error message

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

What it means

PytorchTypeConstraint.type_check verifies that a torch.Tensor batch has the dtype declared in the PytorchTensor[dtype, shape] hint. When tensor.dtype differs from the constraint's dtype, Beam raises this TypeError at runtime, because downstream code (e.g. a model) expects a specific dtype such as float32.

Solutions

1. Cast the tensor to the hinted dtype before yielding it: batch.to(torch.float32) (or .float()/.double() as appropriate).
2. Create tensors with the intended dtype explicitly: torch.tensor(data, dtype=torch.float32).
3. Change the PytorchTensor hint's dtype to match the actual runtime dtype, e.g. PytorchTensor[torch.float64, (N, 128)].
4. Ensure numpy conversions use matching dtypes (arr.astype(np.float32) before torch.from_numpy).

Example fix

// before
yield torch.from_numpy(embeddings)  # float64 from numpy, hint is float32
// after
yield torch.from_numpy(embeddings.astype(np.float32))  # matches PytorchTensor[torch.float32, ...]

def ensure_dtype(t: torch.Tensor, dtype) -> torch.Tensor:
    return t if t.dtype == dtype else t.to(dtype)

def has_dtype(t, dtype) -> bool:
    return isinstance(t, torch.Tensor) and t.dtype == dtype

try:
    yield batch
except TypeError as e:
    if 'does not have expected dtype' in str(e):
        yield batch.to(expected_dtype)
    else:
        raise

When it happens

Trigger: Annotating a PCollection with PytorchTensor[torch.float32, ...] while the DoFn emits tensors created as float64 (e.g. converted from numpy float64 arrays via torch.from_numpy), or after a .double()/.half() cast, when runtime type checking is enabled.

Common situations: torch.from_numpy preserving float64 numpy dtype, mixing a half-precision model with float32 hints, loading checkpoints saved in a different dtype, or torch.tensor on integer data producing int64 where a float hint was declared.

Related errors

• Batch does not have expected shape
• Batch is not an instance of torch.Tensor
• batch type and element type must have equivalent dtypes
• A has been supplied to the model handler, but the required…
• Batch does not have expected dtype

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

Appendix: source

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

  def get_length(self, batch):
    return batch.size(dim=self.partition_dimension)

  def estimate_byte_size(self, batch):
    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()

View on GitHub (pinned to 12126d8942)