def METHOD_NAME(carry, x):
    """Convert JAX scan arguments into format expected by scan_inner_func.
    scan(carry, x) -> scan_inner_func(seqs, mit_sot, sit_sot, shared, non_seqs)
    """
    # `carry` contains all inner taps, shared terms, and non_seqs
    (
        inner_mit_sot,
        inner_sit_sot,
        inner_shared,
        inner_non_seqs,
    ) = carry
    # `x` contains the inner sequences
    inner_seqs = x
    mit_sot_flatten = []
    for array, index in zip(inner_mit_sot, op.info.mit_sot_in_slices):
        mit_sot_flatten.extend(array[jnp.array(index)])
    inner_scan_inputs = [
        *inner_seqs,
        *mit_sot_flatten,
        *inner_sit_sot,
        *inner_shared,
        *inner_non_seqs,
    ]
    return inner_scan_inputs