Trying to use the tensorflow.plugins.hparams in tensorflow 2.0 to create a bunch of different optimizers

I am trying to tune hyperparameters for a neural network using tensorflow.plugins.hparams.

In this link, it gives a suggested code of how to use the function to tune hyperparameters.

As is visible in the provided link, one could use the following:

HP_NUM_UNITS = hp.HParam('num_units', hp.Discrete([16, 32]))
HP_DROPOUT = hp.HParam('dropout', hp.RealInterval(0.1, 0.2))
HP_OPTIMIZER = hp.HParam('optimizer', hp.Discrete(['adam', 'sgd']))

METRIC_ACCURACY = 'accuracy'

with tf.summary.create_file_writer('logs/hparam_tuning').as_default():
  hp.hparams_config(
    hparams=[HP_NUM_UNITS, HP_DROPOUT, HP_OPTIMIZER],
    metrics=[hp.Metric(METRIC_ACCURACY, display_name='Accuracy')],
  )

def train_test_model(hparams):
  model = tf.keras.models.Sequential([
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(hparams[HP_NUM_UNITS], activation=tf.nn.relu),
    tf.keras.layers.Dropout(hparams[HP_DROPOUT]),
    tf.keras.layers.Dense(10, activation=tf.nn.softmax),
  ])
  model.compile(
      optimizer=hparams[HP_OPTIMIZER],
      loss='sparse_categorical_crossentropy',
      metrics=['accuracy'],
  )

  model.fit(x_train, y_train, epochs=1) # Run with 1 epoch to speed things up for demo purposes
  _, accuracy = model.evaluate(x_test, y_test)
  return accuracy

I'll focus on this next line as I took it as a reference to what I want to do. The line:

tf.keras.layers.Dense(hparams[HP_NUM_UNITS], activation=tf.nn.relu),

What I want to eventually do is create a bunch of different types of optimizers with different learning rates, decaying rates, beta_1 and beta_2 values, etc. This is what I tried to do:

HP_LR = hp.HParam('learning_rate',hp.Discrete([0.01,0.005,0.001,0.0005,0.0001,0.00005,0.00001,0.000005,0.000001,0.0000005,0.0000001]))
HP_MOM1 = hp.HParam('momentum_beta1',hp.Discrete([0.9,0.99,0.999,0.9999,0.99999]))
HP_MOM2 = hp.HParam('beta2',hp.Discrete([0.99,0.999,0.9999,0.99999,0.999999]))
HP_BOOL = hp.HParam('amsgrad',hp.Discrete(['True','False']))

METRIC_LOSS = 'dl_tf_loss'

with tf.summary.create_file_writer('logs/hparam_tuning').as_default():
  hp.hparams_config(hparams=[HP_LR, HP_MOM1, HP_MOM2, HP_BOOL],metrics=[hp.Metric(METRIC_LOSS, display_name='Loss')],)

That runs fine but then I try to do

sgd=keras.optimizers.SGD(hparams[HP_LR], decay=0, momentum=hparams[HP_MOM1], nesterov=hparams[HP_BOOL])

and I get the following error:

KeyError                                  Traceback (most recent call last)
<ipython-input-53-1ff09e6440a9> in <module>()
----> 1 sgd=keras.optimizers.SGD(hparams[HP_LR], decay=0, momentum=hparams[HP_MOM1], nesterov=hparams[HP_BOOL])
      2 #RMSprop=keras.optimizers.RMSprop(lr=hparams[HP_LR], rho=0.9, epsilon=None, decay=0.0)
      3 #adam=keras.optimizers.Adam(lr=hparams[HP_LR], beta_1=hparams[HP_MOM1], beta_2=hparams[HP_MOM2], epsilon=None, decay=0.0, amsgrad=hparams[HP_BOOL])
      4 
      5 #HP_OPTIMIZER = hp.HParam('optimizers', hp.Discrete([sgd, RMSprop,adam]))

KeyError: HParam(name='learning_rate', domain=Discrete([1e-07, 5e-07, 1e-06, 5e-06, 1e-05, 5e-05, 0.0001, 0.0005, 0.001, 0.005, 0.01]), display_name=None, description=None)

I would really appreciate it if someone could please let me understand how I can create a bunch of different optimizers using different sets of learning rates, decays, etc.

Thanks!