def _create_slots(self, var_list):
super()._create_slots(var_list)
for var in var_list:
try:
sharding = gshard_utils.GetVarSharding(var)
except ValueError:
continue
if sharding.is_replicated:
continue
m = self.get_slot(var, 'm')
v = self.get_slot(var, 'v')
sharding.ApplyToVariable(m)
sharding.ApplyToVariable(v)
def _create_slots(self, var_list):
if not self._counter:
self._counter = tf.get_variable(
shape=[], initializer=tf.zeros_initializer, name='update_count')
for v in var_list:
vo = self._opt._zeros_slot(v, 'grad_accum', 'GradientAccumulator') # pylint: disable=protected-access
sharding = None
try:
sharding = gshard_utils.GetVarSharding(v)
except ValueError:
continue
if sharding and not sharding.is_replicated:
sharding.ApplyToVariable(vo)