[Bug] AutoEnsembleEstimator() cannot be instatiated if tf< 2.x

[le-dawg]

Intent:
I am following the code sample given by @cweill and given in the docs v0.8.0 to instatiate an AutoEnsembleEstimator() with the simplest use case of one candidate.

Error:
The given object is not an Optimizer instance. Given: <tensorflow.python.keras.optimizer_v2.rmsprop.RMSprop object at 0x7fd6f61ca160>
Minimal reproducible example HERE
Call context:

candmax_iteration_steps = TRAIN_STEPS // ADANET_ITERATIONS

# Learn to ensemble linear and DNN models.
adaestimator = adanet.AutoEnsembleEstimator(
    head=head,
    candidate_pool=lambda config: {
        "linearest":
            tf.estimator.LinearEstimator(
                head=head,
                feature_columns=feature_columns,
                optimizer = lambda: tf.compat.v2.optimizers.RMSprop(),
                config=make_config("ada_linearest"))},max_iteration_steps=candmax_iteration_steps)

Details:
The same error occurs when

• no optimizer argument is given (AutoEnsembleEstimator() defaults to FtrlOptimizer by inheritance)
• any other optimizer is given using any other convention
  • ... tf.compat.v1.keras.optimizers....
  • ... tf.compat.v2.keras.optimizers...

The error, possibly, lies within .../tensorflow_estimator/python/estimator/head/base_head.py.
The check for compatibility appears broken. Not sure if this is still an issue within AdaNet. lambda trick

[cweill]

@le-dawg We just released adanet v0.9.0 which drops TF 1.* support. Could you please try upgrading adanet to the latest version?