def convert_pixel_to_2d_indices(edge_length, flattened_pixel_index):
  """Maps an index of an array to its reshaped 2D matrix's rows and columns.

  This function maps the index of an array with length edge_length ** 2 to the
  rows and columns of its reshaped 2D matrix with shape
  (edge_length, edge_length).

  Args:
    edge_length: int, side length of the 2D array (image) whose pixels are to be
        masked.
    flattened_pixel_index: int, flattened pixel index in the image in
        a row major order.
  Returns:
    row_index: int, row index of the 2D array
    column_index: int, column index of the 2D array
  """
  return (
      flattened_pixel_index // edge_length, flattened_pixel_index % edge_length)