Unable to set selection for DataTable with custom action settings #3030

dash                 2.18.1
dash-core-components 2.0.0
dash-html-components 2.0.0
dash-table           5.0.0

import dash
from dash.dependencies import Input, Output, State
from dash import dcc, html
from dash import dash_table as dt

app = dash.Dash(__name__,
                suppress_callback_exceptions=True)

app.layout = html.Div([
    html.Div(dt.DataTable(
        id='datatable',
        columns=[
            {'name': 'a', 'id': 'a'},
            {'name': 'b', 'id': 'b'},
        ],
        data=[
            {'id': 0, 'a': 1, 'b': 4},
            {'id': 1, 'a': 2, 'b': 5},
            {'id': 2, 'a': 3, 'b': 6},
        ],
        editable=True,
        filter_action='native',
        sort_action='native',
        sort_by=[],
        sort_mode='multi',
        row_selectable='multi',
        selected_rows=[0, 1],
        page_action='native',
        page_current= 0,
        page_size= 10,
    ), id='datatable-div'),
    dcc.RadioItems(['Native', 'Custom'], 'Native', id='table-mode', style={'color': 'white'}),
    html.Button('Set Selection', id='set-selection', style={'color': 'white'}),
])

@app.callback(
    [Output('datatable', 'selected_rows'),
     Output('datatable', 'sort_by'),],
    [Input('set-selection', 'n_clicks')],
    [State('datatable', 'sort_by')], 
    prevent_initial_call=True)
def set_selection(clicks, sort_by):
    return [
        [0, 1, 2],
        sort_by,
    ]

@app.callback(
    [Output('datatable', 'filter_action'),
     Output('datatable', 'sort_action'),
     Output('datatable', 'page_action')],
    [Input('table-mode', 'value')], 
    prevent_initial_call=True)
def set_table_mode(value):
    return [value.lower()] * 3

    
if __name__ == '__main__':
    app.run(debug=True)