Add support for reshape
Fix test case