def METHOD_NAME(self, layer_cls, init_args):
    num_classes = 10
    layer = layer_cls(**init_args)
    img = tf.random.uniform(
        shape=(3, 512, 512, 3), minval=0, maxval=1, dtype=tf.float32
    )
    segmentation_masks = tf.random.uniform(
        shape=(3, 512, 512, 1), minval=0, maxval=num_classes, dtype=tf.int32
    )
    inputs = {"images": img, "segmentation_masks": segmentation_masks}
    outputs = layer(inputs)
    self.assertIn("segmentation_masks", outputs)
    # This currently asserts that all layers are no-ops.
    # When preprocessing layers are updated to mutate segmentation masks,
    # this condition should only be asserted for no-op layers.
    self.assertAllClose(
        inputs["segmentation_masks"], outputs["segmentation_masks"]
    )