Fix per-feature boundary lost by ZCH index dedup (#4525)

[yingufan]

Summary:

When a managed collision table has more than one feature, index dedup can move
ids from one feature to another.

All features of a table share one hash range (see _create_dedup_indices), so
fbgemm.jagged_unique_indices dedups them together. The op cannot tell which
feature each surviving id came from, so it just splits the total evenly over the
table's (feature, batch) slots. That is fine for a table with one feature. With
two features of very different lengths it is not: a 33 long feature and a 1280
long feature both come back as roughly 656, so ids end up under the wrong
feature.

This is a correctness problem and not just wrong bookkeeping, because ZCH decides
whether to insert an id from the name of the feature it arrived under: a feature
whose name ends in the read-only suffix is looked up but never written. Ids moved
to the wrong feature get written when they should not be.

The fix walks each surviving id back to the first feature it appeared in,
regroups the values so each feature is contiguous again, and recounts. Ids sent
by both features still collapse into one row, so we keep the dedup saving. It
only runs when a table actually has more than one feature, so the common path is
unchanged.

Also declares the reverse_indices field that _dedup_indices already appends
to. mc_modules has its own local EmbeddingCollectionContext that was missing
it, so sharding a ManagedCollisionCollection on its own with dedup enabled
raised AttributeError. It only worked through
ManagedCollisionEmbeddingCollection, whose context comes from embedding.py.

Differential Revision: D115457342

[meta-codesync]

@yingufan has exported this pull request. If you are a Meta employee, you can view the originating Diff in D115457342.