Added the square function