add new imbalance module store, which allows to store weights in an atom property
(cherry picked from commit 5405622f3bdacf79e99304fedf473df013656f92)