def tone_mapper(image, gamma):
    """Applies gamma correction to the input image.

    Tone maps the input image batch in order to make scenes with a high dynamic
    range viewable. The gamma correction factor is computed separately per image,
    but is shared between all provided channels. The exact function computed is:

    image_out = A*image_in^gamma, where A is an image-wide constant computed so
    that the maximum image value is approximately 1. The correction is applied
    to all channels.

    Args:
        image: 4-D float32 tensor with shape [batch_size, image_height,
            image_width, channel_count]. The batch of images to tone map.
        gamma: 0-D float32 nonnegative tensor. Values of gamma below one compress
            relative contrast in the image, and values above one increase it. A
            value of 1 is equivalent to scaling the image to have a maximum value
            of 1.
    Returns:
        4-D float32 tensor with shape [batch_size, image_height, image_width,
        channel_count]. Contains the gamma-corrected images, clipped to the range
        [0, 1].
    """
    batch_size = image.shape[0]
    corrected_image = image ** gamma
    image_max, _ = corrected_image.view(batch_size, -1).max(dim=1)
    scaled_image = corrected_image / image_max.view(batch_size, 1, 1, 1)
    return scaled_image.clamp(0.0, 1.0)