import torch


def point_edt2(point: torch.Tensor, grid: torch.Tensor) -> torch.Tensor:
    """Batched version of a Squared Euclidean Distance Transform for a D-dimensional point.

    Args:
      point: torch.Tensor of size [batch, *, D]. Each element is interpreted as a D-dimensional point (e.g. [row, col]
        or [depth, row, col]) and * represents any number of additional dimensions (eg. channels or images or both).
      grid: torch.Tensor of size [*D, D], where *D represents D elements defining a lattice that defines the coordinates
        of each output pixel/voxel.

    Returns:
      a torch.Tensor of size [batch, *, *D] representing the EDT^2 of each point in the input batch where *
      represents any additional dimensions from the input, and *D is the size of each dimension of the lattice.

    """
    inshape = point.shape
    outshape = (*inshape[0:-1], *grid.shape[0:-1])
    dim = len(grid.shape) - 1
    point = point.view(-1, *[1] * dim, dim)

    # need to replicate the grid for each item in the batch
    grid = grid.expand(point.shape[0], *grid.shape)

    pl = (grid - point)
    d = (pl * pl).sum(dim=-1)

    d = d.view(outshape)

    return d