def _dct_8x8(input: Tensor) -> Tensor:
    """Perform an 8 x 8 discrete cosine transform.

    Args:
        input (Tensor): Patched input tensor of the shape :math:`(B, N, 8, 8)`.

    Returns:
        output (Tensor): DCT output tensor of the shape :math:`(B, N, 8, 8)`.

    """
    # Get dtype and device
    dtype: Dtype = input.dtype
    device: Device = input.device
    # Make DCT tensor and scaling
    index: Tensor = torch.arange(8, dtype=dtype, device=device)
    x, y, u, v = torch.meshgrid(index, index, index, index)
    dct_tensor: Tensor = ((2.0 * x + 1.0) * u * pi / 16.0).cos() * ((2.0 * y + 1.0) * v * pi / 16.0).cos()
    alpha: Tensor = torch.ones(8, dtype=dtype, device=device)
    alpha[0] = 1.0 / (2**0.5)
    dct_scale: Tensor = torch.einsum("i, j -> ij", alpha, alpha) * 0.25
    # Apply DCT
    output: Tensor = dct_scale[None, None] * torch.tensordot(input - 128.0, dct_tensor)
    return output
Perform an 8 x 8 discrete cosine transform.

    Args:
        input (Tensor): Patched input tensor of the shape :math:`(B, N, 8, 8)`.

    Returns:
        output (Tensor): DCT output tensor of the shape :math:`(B, N, 8, 8)`.

    
kornia/enhance/jpeg.py