NumPy shapes can be int or long, adjust assertions to accomodate