Just curious, what if when it iterate through the param groups, it uses sparse_adam when grads are sparse and dense adam when they are not? Then just adding sparse=True on a unit would automatically switch that unit to sparse_adam.
As it is now, we have to go through and find all the param groups that we know are going to be sparse, and then pass them to a separate optimizer.