Maintaining Data Color After Dropdown Updates In Plotly

I'm trying to create a plot using Plotly that allows you to select from dropdown menus what features are being plotted on the x and y axis. My approach works, but there's a set of actions that remove the coloring of the points being plotted.

Here's a Colab with the steps to reproduce this written out, and done with minimal code (Plotly plays nice with Colab): https://colab.research.google.com/drive/19PCS8QH9n6VVN9UBOKMay99VuSXq1QGG?usp=sharing

If you want to use your own environment, the following code will reproduce the issue after you've done the following 2 steps:

1. Pick one of the two dropdown menus and change the selected value at least one time
2. Change the selected value on the dropdown menu you have not changed yet

You should then see that the original coloring of the points is lost.

import numpy as np
import pandas as pd

import plotly.express as px
import plotly.io as pio


pio.templates.default = "plotly_dark"


def get_correlation_figure_please(merged_df):
    cols = [col for col, t in zip(merged_df.columns, merged_df.dtypes) if t != object]

    start_dropdown_indices = [0, 0]

    # Create the scatter plot of the initially selected variables
    fig = px.scatter(
        merged_df,
        x=cols[start_dropdown_indices[0]],
        y=cols[start_dropdown_indices[1]],
        color='serial_number_id',
    )

    # Create the drop-down menus which will be used to choose the desired file characteristics for comparison
    drop_downs = []
    for axis in ['x', 'y']:
        drop_downs.append([
            dict(
                method = 'update', 
                args = [
                    {axis : [merged_df[cols[k]]]},
                    {'%saxis.title.text'%axis: cols[k]},
                    # {'color':[merged_df['serial_number_id']],'color_discrete_map':SERIALS_TO_INDEX},
                ],
                label = cols[k]) for k in range(len(cols))
            ])
      
    # Sets up various apsects of the Plotly figure that is currently being produced.  This ranges from 
    # aethetic things, to setting the dropdown menues as part of the figure
    fig.update_layout(
        title_x=0.4,
        showlegend=False,
        updatemenus=[{
            'active': start_j,
            'buttons': drop_down,
            'x': 1.125,
            'y': y_height,
            'xanchor': 'left',
            'yanchor': 'top',
          } for drop_down, start_j, y_height in zip(drop_downs, start_dropdown_indices, [1, .85])])

    return fig


# Set up a dummy dataframe with 20 points each with 5 featuers
df = pd.DataFrame({str(j):np.random.rand(20) for j in range(5)})

# Set up a column of dummied serial numbers (to be used to decide the coloring of each point)
df['serial_number_id'] = df['1'].map(lambda x : '0' if x < 1/3 else ('1' if x < 2/3 else '2'))


fig = get_correlation_figure_please(df)
fig.show()