import torch


def image_positional_encoding(shape):
  """Generates *per-channel* positional encodings for 2d images.

  The positional encoding is a Tensor of shape (N, 2*C, H, W) of (x, y) pixel 
  coordinates scaled to be between -.5 and .5. 

  Args: 
    shape: NCHW shape of image for which to generate positional encodings.
  Returns:
    The positional encodings.
  """
  n, c, h, w = shape
  zeros = torch.zeros(n, c, h, w) 
  return torch.cat((
    (torch.arange(-.5, .5, 1 / h)[None, None, :, None] + zeros),
    (torch.arange(-.5, .5, 1 / w)[None, None, None, :] + zeros)),
    dim=1)