The TrainableWrapper still not works well with variable_scope + partitioner , We will spend long time to push a PR to TensorFlow core to resolve this problem and as an equivalent alternative, you can use partitioner in tf.get_variable directly.

Adam's _create_slots function is colocate beta1_power and beta2_power the variable with the lowest name in var_list, plz refer to here. When BN is turned off(maybe triggered by some other special conditions), the dynamic_embedding TrainableWrapper in var_list may be brought to the front . When colocate_with TrainableWrapper's the device that is not /job:ps, but /job:worker, resulting in the variable is not initialized. Because we can not change the adam._create_slots in TF and we also had to keep the TrainableWrapper on /job:worker local. Then we recommend the way to avoid this issue is changing the first letter of the name of the tfra.embedding_lookup(name='a-to-z-xxxx') to ‘z’ so that it will be queued to the back. In the furture we will make core to know the existing of tfra.dynamic_embedding.TrainableWrapper as fundamental solution.