import torch
from typing import Optional
from typing import Tuple


def generate_image_coords(
  image_shape: torch.Size,
  dtype: Optional[torch.dtype] = None,
  device: Optional[torch.device] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
  """Generate image horizontal, vertical coordinates

  Args:
    image_shape (torch.Size): shape of image. Expected 2/3/4D or higher
    dimensions, where the last two dimensions are (h, w).
    flip_h (bool, optional): whether to flip horizontal axis. Defaults to True.
    dtype (torch.dtype, optional): data type, defaults to torch.float32.
      Defaults to None.
    device (torch.device, optional): torch device. Defaults to None.

  Returns:
    torch.tensor: horizontal coordinates in `image_shape`.
    torhc.tensor: vertical coordinates in `image_shape`.
  """
  dtype = dtype or torch.float32
  ndims = len(image_shape)
  if ndims < 2:
    raise ValueError("rank of `image_shape` must be at east 2D, "
      f"got {ndims}")
  h = image_shape[-2]
  w = image_shape[-1]
  # Generate x, y coordinates
  x = torch.arange(w, dtype=dtype, device=device)
  y = torch.arange(h, dtype=dtype, device=device)
  # Expand dims to match depth map
  x = x.view((1,)*(ndims-2) + (1, -1)) # (..., 1, w)
  y = y.view((1,)*(ndims-2) + (-1, 1)) # (..., h, 1)
  x = torch.broadcast_to(x, image_shape) # (..., h, w)
  y = torch.broadcast_to(y, image_shape)
  return x, y