cross-batch runs not working about fastscape HOT 6 CLOSED

with basic_model, xs.monitoring.ProgressBar():
    ds_out = (
        ds_in
        .stack(batch=['U','K'])
        .drop_indexes(["batch", "U", "K"])    # explicitly drop the multi-index 
        .drop_vars("batch")                   # remove the batch dimension coordinate
        .xsimlab.run(
            check_dims='transpose',
            store=zgroup,
            batch_dim='batch',
            parallel=True,
            scheduler='processes',
        )
        .set_index(batch=["U", "K"])          # explicitly set the multi-index
        .unstack("batch")
    )

ds_in = xs.create_setup(
    input_vars={
        'uplift__rate': ('uplift__rate', [1e-4,1e-3]),
        'spl__k_coef': ('spl__k_coef', [1e-6,1e-5]),
        ...
    },
)

with basic_model, xs.monitoring.ProgressBar():
    ds_out = (
        ds_in
        .stack(batch=['uplift__rate', 'spl__k_coef'], create_index=False)    # do not create an index
        .xsimlab.run(
            check_dims='transpose',
            store=zgroup,
            batch_dim='batch',
            parallel=True,
            scheduler='processes',
        )
        .set_index(batch=["uplift__rate", "spl__k_coef"])           # explicitly set the multi-index
        .unstack("batch")
    )