[IMP-dev] Add square function