"""Module implements DEPT Embeddings modules."""
# TODO(anon, anon): Create an embedding module with a cpu full copy and gpu trimmed
# one that depends on a list of indices to keep. Implement reset functions and remap
# functions (back and forth the GPU one). Refs: https://github.com/anon/private_mt_dao/blob/b038d1d69571408c364dc3c172caa6b85cc9e27c/flower_llm/strategy/aggregation.py#L250
