def posterize(input: Tensor, bits: Union[int, Tensor]) -> Tensor:
    r"""Reduce the number of bits for each color channel.

    .. image:: _static/img/posterize.png

    Non-differentiable function, ``torch.uint8`` involved.

    Args:
        input: image tensor with shape :math:`(*, C, H, W)` to posterize.
        bits: number of high bits. Must be in range [0, 8].
            If int or one element tensor, input will be posterized by this bits.
            If 1-d tensor, input will be posterized element-wisely, len(bits) == input.shape[-3].
            If n-d tensor, input will be posterized element-channel-wisely, bits.shape == input.shape[:len(bits.shape)]

    Returns:
        Image with reduced color channels with shape :math:`(*, C, H, W)`.

    Example:
        >>> x = torch.rand(1, 6, 3, 3)
        >>> out = posterize(x, bits=8)
        >>> torch.testing.assert_close(x, out)

        >>> x = torch.rand(2, 6, 3, 3)
        >>> bits = torch.tensor([4, 2])
        >>> posterize(x, bits).shape
        torch.Size([2, 6, 3, 3])

    """
    if not isinstance(input, Tensor):
        raise TypeError(f"Input type is not a Tensor. Got {type(input)}")

    if not isinstance(bits, (int, Tensor)):
        raise TypeError(f"bits type is not an int or Tensor. Got {type(bits)}")

    if isinstance(bits, int):
        bits = torch.as_tensor(bits)

    # TODO: find a better way to check boundaries on tensors
    # if not torch.all((bits >= 0) * (bits <= 8)) and bits.dtype == torch.int:
    #     raise ValueError(f"bits must be integers within range [0, 8]. Got {bits}.")

    # TODO: Make a differentiable version
    # Current version:
    # Ref: https://github.com/open-mmlab/mmcv/pull/132/files#diff-309c9320c7f71bedffe89a70ccff7f3bR19
    # Ref: https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/autoaugment.py#L222
    # Potential approach: implementing kornia.LUT with floating points
    # https://github.com/albumentations-team/albumentations/blob/master/albumentations/augmentations/functional.py#L472
    def _left_shift(input: Tensor, shift: Tensor) -> Tensor:
        return ((input * 255).to(torch.uint8) * (2**shift)).to(input.dtype) / 255.0

    def _right_shift(input: Tensor, shift: Tensor) -> Tensor:
        return (input * 255).to(torch.uint8) / (2**shift).to(input.dtype) / 255.0

    def _posterize_one(input: Tensor, bits: Tensor) -> Tensor:
        # Single bits value condition
        if bits == 0:
            return torch.zeros_like(input)
        if bits == 8:
            return input.clone()
        bits = 8 - bits
        return _left_shift(_right_shift(input, bits), bits)

    if len(bits.shape) == 0 or (len(bits.shape) == 1 and len(bits) == 1):
        return _posterize_one(input, bits)

    res = []
    if len(bits.shape) == 1:
        if bits.shape[0] != input.shape[0]:
            raise AssertionError(
                f"Batch size must be equal between bits and input. Got {bits.shape[0]}, {input.shape[0]}."
            )

        for i in range(input.shape[0]):
            res.append(_posterize_one(input[i], bits[i]))
        return torch.stack(res, dim=0)

    if bits.shape != input.shape[: len(bits.shape)]:
        raise AssertionError(
            "Batch and channel must be equal between bits and input. "
            f"Got {bits.shape}, {input.shape[: len(bits.shape)]}."
        )
    _input = input.view(-1, *input.shape[len(bits.shape) :])
    _bits = bits.flatten()
    for i in range(input.shape[0]):
        res.append(_posterize_one(_input[i], _bits[i]))
    return torch.stack(res, dim=0).reshape(*input.shape)
Reduce the number of bits for each color channel.

    .. image:: _static/img/posterize.png

    Non-differentiable function, ``torch.uint8`` involved.

    Args:
        input: image tensor with shape :math:`(*, C, H, W)` to posterize.
        bits: number of high bits. Must be in range [0, 8].
            If int or one element tensor, input will be posterized by this bits.
            If 1-d tensor, input will be posterized element-wisely, len(bits) == input.shape[-3].
            If n-d tensor, input will be posterized element-channel-wisely, bits.shape == input.shape[:len(bits.shape)]

    Returns:
        Image with reduced color channels with shape :math:`(*, C, H, W)`.

    Example:
        >>> x = torch.rand(1, 6, 3, 3)
        >>> out = posterize(x, bits=8)
        >>> torch.testing.assert_close(x, out)

        >>> x = torch.rand(2, 6, 3, 3)
        >>> bits = torch.tensor([4, 2])
        >>> posterize(x, bits).shape
        torch.Size([2, 6, 3, 3])

    
kornia/enhance/adjust.py