As your code is not a reproducible example, it is hard to give guidance. However, I think you should consider Ordering of cells via mesh.topology.original_cell_index - #2 by dokken using the original cell index map to transfer the entity tags from its input index.