keras/src/ops/numpy.py/ExpandDims/__init__
class ExpandDims:    def __init__(self, axis):
        super().__init__()
        if not isinstance(axis, (int, tuple, list)):
            raise ValueError(
                "The `axis` argument to `expand_dims` should be an integer, "
                f"tuple or list. Received axis={axis}"
            )
        self.axis = axis