Dev update: Better fix for flatten()

- Smaller code and works with a mix of scalars and arrays