input: Write a python function to generate an identity matrix of size n x n.
  assert generate_matrix(3) == [[1, 0, 0], [0, 1, 0], [0, 0, 1]]
  assert generate_matrix(4) == [[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]]
  assert generate_matrix(2) == [[1, 0], [0, 1]]
output: def generate_matrix(n):
      return [[1 if i == j else 0 for j in range(n)] for i in range(n)]