use more accurate reshape
add additional tests for different input cases