Can you use assert_array_equal here. That np.all call does not check the shape.