def METHOD_NAME(self):
    num_classes = 8
    input_images = np.random.random((2, 20, 20, 3)).astype(np.float32)
    # Masks are all 0s or 8s, to verify that when we rotate we don't do bad
    # mask interpolation to either a 0 or a 7
    masks = np.random.randint(2, size=(2, 20, 20, 1)) * (num_classes - 1)
    inputs = {"images": input_images, "segmentation_masks": masks}
    # Attempting to rotate a sparse mask without specifying num_classes
    # fails.
    bad_layer = RandomRotation(factor=(0.25, 0.25))
    with self.assertRaisesRegex(ValueError, "masks must be one-hot"):
        outputs = bad_layer(inputs)
    # 90 degree rotation.
    layer = RandomRotation(
        factor=(0.25, 0.25), segmentation_classes=num_classes
    )
    outputs = layer(inputs)
    expected_masks = np.rot90(masks, axes=(1, 2))
    self.assertAllClose(expected_masks, outputs["segmentation_masks"])
    # 45-degree rotation. Only verifies that no interpolation takes place.
    layer = RandomRotation(
        factor=(0.125, 0.125), segmentation_classes=num_classes
    )
    outputs = layer(inputs)
    self.assertAllInSet(outputs["segmentation_masks"], [0, 7])