from typing import List
import torch


def compute_grid(shape: List[int]) -> torch.Tensor:
    """
    Computes the normalized and centered grid coordinates for an image
    :param shape: contains batch_size, height and width of grid
    :return: shape (B, 2, H, W) containing grid coordinates
    """
    batch_size, height, width = shape

    h_range = torch.arange(height)
    w_range = torch.arange(width)
    grid_x = w_range.repeat(height, 1).unsqueeze(0)
    grid_y = h_range.view(-1, 1).repeat(1, width).unsqueeze(0)

    return torch.cat(
        (2 * grid_x - width + 1, height - 2 * grid_y - 1), 0
    ).repeat(batch_size, 1, 1, 1)