I prefer the current alignment (close brace under open brace), but I agree that it's a little fuddled. Maybe `segments_color],\n axis=-1)`?