def check_shape(
array: Union[np.ndarray, torch.Tensor],
expected_shape: Union[Tuple[int, ...], List[Tuple[int, ...]]]
) -> None:
"""
Validates that an input array matches one or more expected shapes.
Dimensions may be integers or `-1`, where `-1` acts as a wildcard.
Parameters
----------
array : Union[np.ndarray, torch.Tensor]
Input array to validate.
expected_shape : Union[Tuple[int, ...], List[Tuple[int, ...]]]
A single expected shape or a list of acceptable shapes.
Use `-1` or `None` to indicate a wildcard dimension.
Examples
--------
>>> x = np.zeros((4, 8))
>>> check_shape(x, (4, 8)) # exact match
>>> check_shape(x, (-1, 8)) # wildcard batch dimension
>>> check_shape(x, [(-1, 4), (-1, 8)]) # multiple acceptable shapes
"""
# Type check
if not isinstance(array, (np.ndarray, torch.Tensor)):
logger.error(
f"Invalid type: expected `np.ndarray` or `torch.Tensor`, "
f"got `{type(array).__name__}`", TypeError
)
actual_shape = tuple(array.shape)
if isinstance(expected_shape, tuple):
expected_shapes = [expected_shape]
else:
expected_shapes = expected_shape
# Try all possible shapes
for shape in expected_shapes:
if len(shape) != len(actual_shape):
continue
match = True
for a, e in zip(actual_shape, shape):
if e not in (-1, None) and a != e:
match = False
break
# Valid shape found
if match:
return
# If no match found
if isinstance(expected_shape, tuple):
logger.error(
f"Invalid shape: expected `{expected_shape}`, "
f"got `{actual_shape}`", ValueError
)
else:
logger.error(
f"expected one of `{expected_shapes}`, "
f"got `{actual_shape}`", ValueError
)