keras/src/tree/optree_impl.py/_packed_nest_with_indices
def _packed_nest_with_indices(structure, flat, index, sequence_fn=None):
    packed = []
    sequence_fn = sequence_fn or _sequence_like
    for s in _yield_value(structure):
        if is_nested(s):
            new_index, child = _packed_nest_with_indices(
                s, flat, index, sequence_fn
            )
            packed.append(sequence_fn(s, child))
            index = new_index
        else:
            packed.append(flat[index])
            index += 1
    return index, packed