import torch


def generate_location_features(img_dimensions,
                               device,
                               feature_type='int'):
  """Calculates location features for an image.

  This function generates location features for an image. The 2-D
  location features range from -1 to 1 for y and x axes each.

  Args:
    img_dimensions: A list of 2 integers for image's y and x dimension.
    feature_type: The data type of location features, integer or float.

  Returns:
    A 3-D float/long tensor with shape
    `[img_dimension[0], img_dimension[1], 2]`.

  Raises:
    ValueError: Type of location features is neither 'int' nor 'float'.
  """
  if feature_type == 'int':
    y_features = torch.arange(img_dimensions[0], device=device)
    x_features = torch.arange(img_dimensions[1], device=device)
  elif feature_type == 'float':
    y_features = torch.linspace(0, 1, img_dimensions[0], device=device)
    x_features = torch.linspace(0, 1, img_dimensions[1], device=device)
  else:
    raise ValueError(
        'Type of location features should be either int or float.')

  y_features, x_features = torch.meshgrid(y_features, x_features)

  location_features = torch.stack([y_features, x_features], dim=2)

  return location_features