def _get_slot_variable_names(scope_name, var_name, optimization_parameters):
"""Return embedding variable names which are consistent with CPU runs."""
if scope_name:
scope_name = scope_name + '/'
if isinstance(optimization_parameters,
tf.compat.v1.tpu.experimental.AdagradParameters):
return tpu_embedding.AdagradSlotVariableName('{}{}/Adagrad'.format(
scope_name, var_name))
elif isinstance(optimization_parameters,
tf.compat.v1.tpu.experimental.AdamParameters):
return tpu_embedding.AdamSlotVariableNames(
'{}{}/Adam/m'.format(scope_name, var_name),
'{}{}/Adam/v'.format(scope_name, var_name))
elif isinstance(optimization_parameters,
tf.compat.v1.tpu.experimental.FtrlParameters):
return tpu_embedding.FtrlSlotVariableName(
'{}{}/Ftrl'.format(scope_name, var_name), # accumulator
'{}{}/Ftrl_1'.format(scope_name, var_name)) # linear
elif isinstance(
optimization_parameters,
tf.compat.v1.tpu.experimental.StochasticGradientDescentParameters):
return None
else:
raise ValueError(
'Support to infer full variable name '
'for optimization_parameter {} has not been added.'.format(
optimization_parameters))
def _get_slot_variable_names(scope_name, var_name, optimization_parameters):
"""Return embedding variable names which are consistent with CPU runs."""
if isinstance(optimization_parameters, tpu_embedding.AdagradParameters):
return tpu_embedding.AdagradSlotVariableName('{}/{}/Adagrad'.format(
scope_name, var_name))
elif isinstance(optimization_parameters, tpu_embedding.AdamParameters):
return tpu_embedding.AdamSlotVariableNames(
'{}/{}/Adam/m'.format(scope_name, var_name),
'{}/{}/Adam/v'.format(scope_name, var_name))
elif isinstance(optimization_parameters,
tpu_embedding.StochasticGradientDescentParameters):
return None
else:
raise ValueError(
'Support to infer full variable name '
'for optimization_parameter {} has not been added.'.format(
optimization_parameters))
