input: Write a python function to return an array with the squares of the elements in the given array.
  assert square_array([1, 2, 3]) == [1, 4, 9]
  assert square_array([4, 5, 6]) == [16, 25, 36]
  assert square_array([0, -1, -2]) == [0, 1, 4]
output: def square_array(arr):
      return [x**2 for x in arr]