TypeError when tuning keras LSTM hyperparameters with Ax

Preface: I'm a bit of an amateur with python and a bit more of one with tensorflow, but I know enough that I can learn from a high level answer. Types are the only real difficulty as I've had Java beaten into me for the last few months

I'm working on an LSTM for sentiment analysis and I decided to choose the hyperparameters using the Ax framework. From what I've gathered, the process is:

I create an Ax experiment in which a model is created using pre-generated params and trained on a sample of the data. This process is repeated a number of times, after which the optimal param values and learning rate are chosen and can be used to train a model with the full dataset.

When I run the code, a single model is created and trained. Upon hitting ax.complete_trial(), I receive the error

TypeError: type of val must be one of (Dict[str, Union[float, numpy.floating, numpy.integer, Tuple[Union[float, numpy.floating, numpy.integer], Union[float, numpy.floating, numpy.integer, NoneType]]]], float, numpy.floating, numpy.integer, Tuple[Union[float, numpy.floating, numpy.integer], Union[float, numpy.floating, numpy.integer, NoneType]], List[Tuple[Dict[str, Union[str, bool, float, int, NoneType]], Dict[str, Union[float, numpy.floating, numpy.integer, Tuple[Union[float, numpy.floating, numpy.integer], Union[float, numpy.floating, numpy.integer, NoneType]]]]]], List[Tuple[Dict[str, Hashable], Dict[str, Union[float, numpy.floating, numpy.integer, Tuple[Union[float, numpy.floating, numpy.integer], Union[float, numpy.floating, numpy.integer, NoneType]]]]]]); got dict instead

followed by

ValueError: Raw data must be data for a single arm for non batched trials.

According to ax_client.py,

Two custom types used in this class for convenience are `TParamValue` and
    `TParameterization`. Those are shortcuts for `Union[str, bool, float, int]`
    and `Dict[str, Union[str, bool, float, int]]`, respectively.

I attempted to cast the getScore() value into TParamValue to union the parameters, thinking I could then cast the value for raw_data into TParameterization and solve the issue.
This yielded TypeError: Cannot instantiate typing.Union

I've tried different configurations of casting and they all throw one of these two errors. For some reason I can't seem to figure out how to use Union properly and I'm out of ideas.

my code:

def getLstm(out_dim, neurons_dense, embedding_size, max_text_len, dropout_rate):
    inputs = Input(shape=max_text_len)
    x = layers.Embedding(1000, embedding_size)(inputs)
    x = layers.LSTM(units=out_dim)(x)
    x = layers.Dense(neurons_dense, activation='tanh')(x)
    x = layers.Dropout(dropout_rate)(x)

    outputs = layers.Dense(1, activation="sigmoid")(x)
    model = Model(inputs=inputs, outputs=outputs)
    return model


def getScore(p, weight=None):
    max_text_len = p.get('max_text_len')
    backend.clear_session()

    model = getLstm(
        p.get('out_dim'),
        p.get('neurons_dense'),
        p.get('embedding_size'),
        max_text_len,
        p.get('dropout_rate')
    )

    lRate = p.get('learning_rate')
    opt = tf.keras.optimizers.Adam(learning_rate=lRate)
    numEpochs = p.get('num_epochs')
    model.compile(
        optimizer=opt,
        loss=losses.BinaryCrossentropy(),
        metrics=[tf.keras.metrics.AUC()]
    )

    xTrain, xTest, yTrain, yTest = train_test_split(x, y, test_size=0.2)
    xTrain, xVal, yTrain, yVal = train_test_split(xTrain, yTrain,
                                                  test_size=0.12)

    pTrain = pad_sequences(xTrain, maxlen=max_text_len)
    pVal = pad_sequences(xVal, maxlen=max_text_len)

    result = model.fit(x=pTrain[:2000],
                       y=yTrain[:2000],
                       batch_size=p.get('batch_size'),
                       epochs=numEpochs,
                       validation_data=(pVal, yVal)
                       )
    score = np.array(result.history['val_auc'][-1:])
    print(score)
    return score, 0


ax = AxClient()

ax.create_experiment(
    name='lstm-tuning',
    parameters=parameters,
    objective_name='lstm_cv',
    minimize=False
)


def run(p):
    return {'lstm_cv': TParamValue(getScore(p))}


for i in range(50):
    p, idx = ax.get_next_trial()
    ax.complete_trial(trial_index=idx, raw_data=TParameterization(run(p)))

How many English words
do you know?
Test your English vocabulary size, and measure
how many words do you know
Online Test
Powered by Examplum