{"record":{"id":"70f52ec43a3a742b","repo":"apache/beam","slug":"could-not-align-batch-type-s-batch-dimension-with-element","errorCode":null,"errorMessage":"Could not align batch type's batch dimension with element type. (batch type dimensions: {batch_type.shape}, element type dimenstions: {element_type.shape}","messagePattern":"Could not align batch type's batch dimension with element type\\. \\(batch type dimensions: (.+?), element type dimenstions: (.+?)","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"sdks/python/apache_beam/typehints/pytorch_type_compatibility.py","lineNumber":62,"sourceCode":"    if not isinstance(element_type, PytorchTypeHint.PytorchTypeConstraint):\n      element_type = PytorchTensor[element_type, ()]\n\n    if not isinstance(batch_type, PytorchTypeHint.PytorchTypeConstraint):\n      if not batch_type == torch.Tensor:\n        raise TypeError(\n            \"batch type must be torch.Tensor or \"\n            \"beam.typehints.pytorch_type_compatibility.PytorchTensor[..]\")\n      batch_type = PytorchTensor[element_type.dtype, (N, )]\n\n    if not batch_type.dtype == element_type.dtype:\n      raise TypeError(\n          \"batch type and element type must have equivalent dtypes \"\n          f\"(batch={batch_type.dtype}, element={element_type.dtype})\")\n    computed_element_shape = list(batch_type.shape)\n    partition_dimension = computed_element_shape.index(N)\n    computed_element_shape.pop(partition_dimension)\n    if not tuple(computed_element_shape) == element_type.shape:\n      raise TypeError(\n          \"Could not align batch type's batch dimension with element type. \"\n          f\"(batch type dimensions: {batch_type.shape}, element type \"\n          f\"dimenstions: {element_type.shape}\")\n\n    return PytorchBatchConverter(\n        batch_type,\n        element_type,\n        batch_type.dtype,\n        element_type.shape,\n        partition_dimension)\n\n  def produce_batch(self, elements):\n    return torch.stack(elements, dim=self.partition_dimension)\n\n  def explode_batch(self, batch):\n    \"\"\"Convert an instance of B to Generator[E].\"\"\"\n    yield from torch.swapaxes(batch, self.partition_dimension, 0)\n","sourceCodeStart":44,"sourceCodeEnd":80,"githubUrl":"https://github.com/apache/beam/blob/12126d8942aaf848030c478b4c6a28c6af861c66/sdks/python/apache_beam/typehints/pytorch_type_compatibility.py#L44-L80","documentation":"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.","triggerScenarios":"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.","commonSituations":"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 ().","solutions":["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).","If elements are scalars, use batch_type=PytorchTensor[dtype, (N,)] and element_type=PytorchTensor[dtype, ()].","Remove redundant/extra dimensions from the batch hint so after dropping N the tuples match exactly.","Pass plain torch.Tensor as batch_type so Beam constructs PytorchTensor[element.dtype, (N,)] consistent with the element type."],"exampleFix":"// before\nbatch_type=PytorchTensor[torch.float32, (N, 224, 224, 3)],\nelement_type=PytorchTensor[torch.float32, (3, 224, 224)]\n// after\nbatch_type=PytorchTensor[torch.float32, (N, 3, 224, 224)],\nelement_type=PytorchTensor[torch.float32, (3, 224, 224)]","handlingStrategy":"validation","validationCode":"def check_batch_shape(batch_type, element_type):\n    dims = [d for d in batch_type.shape if d != N]\n    if tuple(dims) != tuple(element_type.shape):\n        raise TypeError(f'shape mismatch: batch={batch_type.shape}, element={element_type.shape}')","typeGuard":"def shapes_align(batch_type, element_type) -> bool:\n    dims = [d for d in getattr(batch_type, 'shape', ()) if d != N]\n    return tuple(dims) == tuple(getattr(element_type, 'shape', ()))","tryCatchPattern":null,"preventionTips":["Derive the batch shape mechanically: element shape plus one N at the partition dimension.","Use batch_type=torch.Tensor to auto-generate the (N,) + element shape batch type.","Keep element and batch hints defined together in one constant so they cannot drift.","Test converter registration in CI before launching pipelines."],"tags":["apache-beam","pytorch","typehints","shape"],"backgroundTag":"shape-mismatch","analyzedSha":"12126d8942aaf848030c478b4c6a28c6af861c66","analyzedAt":"2026-09-13T01:50:10.254Z","contentChangedAt":"2026-09-13T01:50:10.254Z","schemaVersion":2},"datasetVersion":"2026-09-20T03:17:13.778Z"}