keras/src/utils/tracking.py/TrackedDict/tree_flatten
class TrackedDict:    def tree_flatten(self):
        from keras.src.utils.module_utils import optree

        # For optree
        keys, values = optree.utils.unzip2(
            optree.utils.total_order_sorted(self.items(), key=lambda kv: kv[0])
        )
        return values, list(keys), keys