def _normalize_and_prepare_optimizer()

in tensorflow_recommenders/layers/embedding/tpu_embedding_layer.py [0:0]

• 25 lines of code
• 8 McCabe index (conditional complexity)

def _normalize_and_prepare_optimizer(optimizer):
  """Normalizes an optimizer into a mid level API optimizer class.

  In the case of a mid level API optimizer, this just passes it through.
  Passing optimizer names, "sgd", "adam", "adagrad" are supported and
  instantiate the mid level API object with default parameters. If a Keras
  optimizer is passed it will be converted to a mid level optimizer.

  Args:
    optimizer: A keras optimizer, string optimizer name or subclass of
  _OptimizationParameters.

  Returns:
    A subclass of tpu_embedding_v2._Optimizer or None.
  """

  if optimizer is None:
    return None
  elif isinstance(optimizer, (tf.tpu.experimental.embedding.SGD,
                              tf.tpu.experimental.embedding.Adagrad,
                              tf.tpu.experimental.embedding.Adam)):
    return optimizer
  elif isinstance(optimizer, str):
    if str(optimizer) == "sgd":
      return tf.tpu.experimental.embedding.SGD()
    elif str(optimizer) == "adagrad":
      return tf.tpu.experimental.embedding.Adagrad()
    elif str(optimizer) == "adam":
      return tf.tpu.experimental.embedding.Adam()
    else:
      raise ValueError("Unknown optimizer name '{}'. Please use one of 'sgd',"
                       "'adagrad' or 'adam'".format(optimizer))
  elif isinstance(optimizer, tf.keras.optimizers.Optimizer):
    return translate_keras_optimizer(optimizer)
  else:
    raise ValueError(
        "Unknown optimizer type {}. Please pass a string optimizer name, a "
        "subclass of keras optimizer or an instance of one of the optimizer "
        "parameter classes under tf.tpu.experimental.embedding.".format(
            type(optimizer)))