def _diamond_square_seed(
    replicates: int,
    width: int,
    height: int,
    random_fn: Callable[..., Tensor],
    device: Optional[torch.device] = None,
    dtype: Optional[torch.dtype] = None,
) -> Tensor:
    """Generate the diamond square image seee.

    Args:
        replicates: the num of batched replicas for the image.
        width: the expected image width.
        height: the expected image height.
        random_fn: the random function to generate the image seed.
        device: the torch device where to create the image seed.
        dtype: the torch dtype where to create the image seed.

    Return:
        the generated image seed of size Bx1xHxW.

    """
    KORNIA_CHECK(width == 3 or height == 3, "Height or Width must be equal to 3.")
    # TODO(anguelos): can we avoid transposing and passing always fixed size. This will cause issues with onnx/jit
    transpose: bool = False
    if height == 3:
        transpose = True
        width, height = height, width

    # width is always 3
    KORNIA_CHECK(height % 2 == 1 and height > 2, "Height must be odd and height bigger than 2")

    res: Tensor = random_fn([replicates, 1, width, height], device=device, dtype=dtype)
    res[..., ::2, ::2] = random_fn([replicates, 1, 2, (height + 1) // 2], device=device, dtype=dtype)

    # Diamond step
    res[..., 1, 1::2] = (res[..., ::2, :-2:2] + res[..., ::2, 2::2]).sum(dim=2) / 4.0

    # Square step
    if width > 3:
        res[..., 1, 2:-3:2] = (
            res[..., 0, 2:-3:2] + res[..., 2, 2:-3:2] + res[..., 1, 0:-4:2] + res[..., 1, 2:-3:2]
        ) / 4.0

    tmp1 = res[..., 2, 0]
    res[..., 1, 0] = res[..., 0, 0] + res[..., 1, 1] + tmp1
    res[..., 1, -1] = res[..., -1, -1] + res[..., 1, -2] + tmp1
    tmp2 = res[..., 1, 1::2]
    res[..., 0, 1::2] = res[..., 0, 0:-2:2] + res[..., 0, 2::2] + tmp2
    res[..., 2, 1::2] = res[..., 2, 0:-2:2] + res[..., 2, 2::2] + tmp2
    res = res / 3.0

    if transpose:
        res = res.transpose(2, 3)
    return res
Generate the diamond square image seee.

    Args:
        replicates: the num of batched replicas for the image.
        width: the expected image width.
        height: the expected image height.
        random_fn: the random function to generate the image seed.
        device: the torch device where to create the image seed.
        dtype: the torch dtype where to create the image seed.

    Return:
        the generated image seed of size Bx1xHxW.

    
kornia/contrib/diamond_square.py