def get_states(self):
        """Get rng states. Copy the dictionary so we have direct
        pointers to the states, not just a pointer to the dictionary."""
        states = {}
        for name in self.states_:
            states[name] = self.states_[name]
        return states
Get rng states. Copy the dictionary so we have direct
        pointers to the states, not just a pointer to the dictionary.
dpkd/transformers/src/transformers/mpu/random.py