improve handling of certain reshape ops through AST transformation

Should fix #6