regression bug (0.5.2->0.6.0): batch wrapped Switch looses it's on_change handler about marimo HOT 2 CLOSED

import marimo

__generated_with = "0.6.0"
app = marimo.App()


@app.cell
def __():
    import marimo as mo
    return mo,


@app.cell
def __(mo):
    get_state, set_state = mo.state(0)
    return get_state, set_state


@app.cell
def __(mo, set_state):
    def on_change(value):
        # don't call `get_state` here!
        set_state(lambda v: v+1)

    switch = mo.ui.switch(value=True, on_change=on_change)
    return on_change, switch


@app.cell
def __(switch):
    switch
    return


@app.cell
def __(mo, switch):
    def wrap_switch(switch):

        wrapped = mo.md("""
        wrapped switch

        {switch}""").batch(switch=switch)
        return wrapped

    wrapped_switch = wrap_switch(switch)
    wrapped_switch
    return wrap_switch, wrapped_switch


@app.cell
def __(wrapped_switch):
    wrapped_switch._elements["switch"]._id
    return


@app.cell
def __(wrapped_switch):
    wrapped_switch.value
    return


@app.cell
def __():
    from marimo._runtime.context import get_context

    get_context().ui_element_registry._objects
    return get_context,


@app.cell
def __(get_state):
    get_state()
    return


if __name__ == "__main__":
    app.run()