def METHOD_NAME(traces):
    """Convert JAX scan traces to PyTensor traces.
    We need to:
        1. Prepend initial states to JAX output traces
        2. Slice final traces if Scan was instructed to only keep a portion
    """
    init_states = mit_sot_init + sit_sot_init + [None] * op.info.n_nit_sot
    buffers = (
        op.outer_mitsot(outer_inputs)
        + op.outer_sitsot(outer_inputs)
        + op.outer_nitsot(outer_inputs)
    )
    partial_traces = []
    for init_state, trace, buffer in zip(init_states, traces, buffers):
        if init_state is not None:
            # MIT-SOT and SIT-SOT: The final output should be as long as the input buffer
            trace = jnp.atleast_1d(trace)
            init_state = jnp.expand_dims(
                init_state, range(trace.ndim - init_state.ndim)
            )
            full_trace = jnp.concatenate([init_state, trace], axis=0)
            buffer_size = buffer.shape[0]
        else:
            # NIT-SOT: Buffer is just the number of entries that should be returned
            full_trace = jnp.atleast_1d(trace)
            buffer_size = buffer
        partial_trace = full_trace[-buffer_size:]
        partial_traces.append(partial_trace)
    return partial_traces