def random_block_mask(
    size: Tuple[int, int, int],
    batch_mask_ratio: float = 0.5,
    min_image_mask_ratio: float = 0.1,
    max_image_mask_ratio: float = 0.5,
    min_num_masks_per_block: int = 4,
    max_num_masks_per_block: Optional[int] = None,
    min_block_aspect_ratio: float = 0.3,
    max_block_aspect_ratio: Optional[float] = None,
    max_attempts_per_block: int = 10,
    device: Optional[Union[torch.device, str]] = None,
) -> Tensor:
    """Creates a random block mask for a batch of images.

    A block is in this context a rectangle of patches in an image that are
    masked together. The function generates block masks until the desired number of
    patches per image are masked. DINOv2 uses a more complex masking strategy that
    only generates masks for mask_ratio of the images. On top of that, it also masks
    a different number of patches for every image. This is controlled by the
    min_image_mask_ratio and max_image_mask_ratio arguments.

    Based on the implementation of the block mask in DINOv2 [0]. For details see [1]
    and [2].

    - [0]: DINOv2, 2023, https://arxiv.org/abs/2304.07193
    - [1]: https://github.com/facebookresearch/dinov2/blob/main/dinov2/data/masking.py
    - [2]: https://github.com/facebookresearch/dinov2/blob/main/dinov2/data/collate.py

    Args:
        size:
            Size of the image batch for which to generate masks.
            Should be (batch_size, height, width).
        batch_mask_ratio:
            Percentage of images per batch for which to generate block masks.
            The remaining images are not masked.
        min_image_mask_ratio:
            Minimum percentage of the image to mask. In practice, fewer than
            min_image_mask_ratio patches of the image can be masked due to additional
            constraints.
        max_image_mask_ratio:
            Maximum percentage of the image to mask.
        min_num_masks_per_block:
            Minimum number of patches to mask per block.
        max_num_masks_per_block:
            Maximum number of patches to mask per block.
        min_block_aspect_ratio:
            Minimum aspect ratio (height/width) of a masked block.
        max_block_aspect_ratio:
            Maximum aspect ratio (height/width) of a masked block.
        max_attempts_per_block:
            Maximum number of attempts to find a valid block mask for an image.
        device:
            Device on which to create the mask.
    Returns:
        A boolean tensor with shape (batch_size, height, width) where each entry
        is True if the patch should be masked and False otherwise.

    Raises:
        ValueError: If 'max_image_mask_ratio' is less than 'min_image_mask_ratio'.
    """

    if max_image_mask_ratio < min_image_mask_ratio:
        raise ValueError(
            "max_image_mask_ratio must be greater or equal to min_image_mask_ratio."
        )

    # B is batch size(number of images), H is height, W is width
    B, H, W = size
    num_images_masked = int(B * batch_mask_ratio)
    probs = torch.linspace(
        min_image_mask_ratio, max_image_mask_ratio, num_images_masked + 1
    ).tolist()
    image_masks = []
    for prob_min, prob_max in zip(probs[:-1], probs[1:]):
        num_mask = int(H * W * random.uniform(prob_min, prob_max))
        image_masks.append(
            random_block_mask_image(
                size=(H, W),
                num_masks=num_mask,
                min_num_masks_per_block=min_num_masks_per_block,
                max_num_masks_per_block=max_num_masks_per_block,
                min_block_aspect_ratio=min_block_aspect_ratio,
                max_block_aspect_ratio=max_block_aspect_ratio,
                max_attempts_per_block=max_attempts_per_block,
                device=device,
            )
        )

    # Add non-masked images to fill the batch
    for _ in range(num_images_masked, B):
        image_masks.append(torch.zeros((H, W), dtype=torch.bool, device=device))

    random.shuffle(image_masks)
    return torch.stack(image_masks)
Creates a random block mask for a batch of images.

    A block is in this context a rectangle of patches in an image that are
    masked together. The function generates block masks until the desired number of
    patches per image are masked. DINOv2 uses a more complex masking strategy that
    only generates masks for mask_ratio of the images. On top of that, it also masks
    a different number of patches for every image. This is controlled by the
    min_image_mask_ratio and max_image_mask_ratio arguments.

    Based on the implementation of the block mask in DINOv2 [0]. For details see [1]
    and [2].

    - [0]: DINOv2, 2023, https://arxiv.org/abs/2304.07193
    - [1]: https://github.com/facebookresearch/dinov2/blob/main/dinov2/data/masking.py
    - [2]: https://github.com/facebookresearch/dinov2/blob/main/dinov2/data/collate.py

    Args:
        size:
            Size of the image batch for which to generate masks.
            Should be (batch_size, height, width).
        batch_mask_ratio:
            Percentage of images per batch for which to generate block masks.
            The remaining images are not masked.
        min_image_mask_ratio:
            Minimum percentage of the image to mask. In practice, fewer than
            min_image_mask_ratio patches of the image can be masked due to additional
            constraints.
        max_image_mask_ratio:
            Maximum percentage of the image to mask.
        min_num_masks_per_block:
            Minimum number of patches to mask per block.
        max_num_masks_per_block:
            Maximum number of patches to mask per block.
        min_block_aspect_ratio:
            Minimum aspect ratio (height/width) of a masked block.
        max_block_aspect_ratio:
            Maximum aspect ratio (height/width) of a masked block.
        max_attempts_per_block:
            Maximum number of attempts to find a valid block mask for an image.
        device:
            Device on which to create the mask.
    Returns:
        A boolean tensor with shape (batch_size, height, width) where each entry
        is True if the patch should be masked and False otherwise.

    Raises:
        ValueError: If 'max_image_mask_ratio' is less than 'min_image_mask_ratio'.
    
lightly/models/utils.py