Unable to hyper parameter tune a model including TensorFlow Probability

[remsforian]

have been working on implementing a Bayesian Neural Network, and wanted to tune the hyperparamiters using the Keras-tuner as I have done in the past with base tensorflow models. I have an example of unlined code that works well:

    def BuildModel(self):
        #we only use an ANN for the pulse profile, so we don't need to split based on the plot
        trainPPdata = self.trainPPdata
        pulse_length = len(trainPPdata[0]) #length of a single data point
        tfd = tfp.distributions

        def prior(kernel_size, bias_size, dtype=None):
            n = kernel_size + bias_size
            prior_model = tf_keras.Sequential(
                [
                    tfp.layers.DistributionLambda(
                        lambda t: tfp.distributions.MultivariateNormalDiag(
                            loc=tf.zeros(n), scale_diag=tf.ones(n)
                        )
                    )
                ]
            )
            return prior_model
        
        
        def posterior(kernel_size, bias_size, dtype=None):
            n = kernel_size + bias_size
            posterior_model = tf_keras.Sequential(
                [
                    tfp.layers.VariableLayer(
                        tfp.layers.MultivariateNormalTriL.params_size(n), dtype=dtype
                    ),
                    tfp.layers.MultivariateNormalTriL(n),
                ]
            )
            return posterior_model
        model = tf_keras.models.Sequential()
        model.add(tf_keras.layers.Dense(units=pulse_length))
        model.add(tfp.layers.DenseVariational(units=64,activation='relu',make_posterior_fn=posterior,make_prior_fn=prior))
        model.add(tf_keras.layers.Dense(units=1,activation='sigmoid'))
        model.compile(
                optimizer=tf_keras.optimizers.Adam(learning_rate = 0.01),
                loss=tf_keras.losses.BinaryCrossentropy(from_logits=False),
                metrics=['accuracy'])
        self.model = model

I also have an example of my attempt at hyperparamiter tuning:

    def Hypertune(self,trials=150,ep_per_trial=10,ex_per_trial=5,pat = 5):
/        trainPPdata = self.trainPPdata
        pulse_length = len(trainPPdata[0]) #length of a single data point
        tfd = tfp.distributions
        
        #creates an array of possible learning rates to be tuned over
        learning_rate_choices = []
        for i in range (-5,0):
            learning_rate_choices.append(10**i)

        activations = ['relu','sigmoid','softmax','softplus','softsign','tanh','selu','elu']
       
        #builds the parts of the model that are to be tuned over

        def prior(kernel_size, bias_size, dtype=None):
            n = kernel_size + bias_size
            prior_model = tf_keras.Sequential(
                [
                    tfp.layers.DistributionLambda(
                        lambda t: tfp.distributions.MultivariateNormalDiag(
                            loc=tf.zeros(n), scale_diag=tf.ones(n)
                        )
                    )
                ]
            )
            return prior_model
        
        
        def posterior(kernel_size, bias_size, dtype=None):
            n = kernel_size + bias_size
            posterior_model = tf_keras.Sequential(
                [
                    tfp.layers.VariableLayer(
                        tfp.layers.MultivariateNormalTriL.params_size(n), dtype=dtype
                    ),
                    tfp.layers.MultivariateNormalTriL(n),
                ]
            )
            return posterior_model


        def build_model(hp):
            hp_learning_rate = hp.Choice('learning rate', values = learning_rate_choices)
            model = tf_keras.models.Sequential()
            model.add(tf_keras.layers.Dense(units=pulse_length))
            for i in range(hp.Int('num layers',0,10)):
                model.add(tfp.layers.DenseVariational(
                    units = hp.Int('units ' + str(i),
                        min_value = 1,
                        max_value = 64,
                        step=1),
                    make_prior_fn = prior, make_posterior_fn = posterior,
                    activation = hp.Choice('activation ' + str(i),activations)))
            model.add(tf_keras.layers.Dense(units=1,activation='sigmoid'))
            model.compile(
                    optimizer=tf_keras.optimizers.Adam(learning_rate = hp_learning_rate),
                    loss=tf_keras.losses.BinaryCrossentropy(from_logits=False),
                    metrics=['accuracy'])
            return model

         tuner = kt.Hyperband(
            build_model,
            objective='val_accuracy',
            max_epochs=ep_per_trial,
            factor=3,
            hyperband_iterations=2) 

        tuner.search(self.trainPPdata,self.trainlabels,epochs=ep_per_trial, validation_split = 0.2, verbose = 1) #does the hypertuning
        self.model = tuner.get_best_models()[0] #saves the tuned model (not trained) 

The traceback I am getting:

Traceback (most recent call last):
  File "/home/tmasters/nonlinear_pacmann/hypertrain.py", line 94, in <module>
    hypertrain_bnn(pulse_profile_bnn)
  File "/home/tmasters/nonlinear_pacmann/hypertrain.py", line 88, in hypertrain_bnn
    bnn.Hypertune(int(config['TUNING']['runs']),int(config['TUNING']['epochs']),int(config['TUNING']['trials']),int(config['TUNING']['patence'])) #hypertunes the paticular cnn (number of tuning runs trials,epochs per tuning trial,trial runs per tune,patence)                                                                                                                                                       
  File "/home/tmasters/nonlinear_pacmann/Classes/bayANN.py", line 159, in Hypertune
    tuner.search(self.trainPPdata,self.trainlabels,epochs=ep_per_trial, validation_split = 0.2, verbose = 1) #does the hypertuning
  File "/home/tmasters/anaconda3/envs/pacmann/lib/python3.10/site-packages/keras_tuner/src/engine/base_tuner.py", line 234, in search
    self._try_run_and_update_trial(trial, *fit_args, **fit_kwargs)
  File "/home/tmasters/anaconda3/envs/pacmann/lib/python3.10/site-packages/keras_tuner/src/engine/base_tuner.py", line 279, in _try_run_and_update_trial
    raise e
  File "/home/tmasters/anaconda3/envs/pacmann/lib/python3.10/site-packages/keras_tuner/src/engine/base_tuner.py", line 274, in _try_run_and_update_trial
    self._run_and_update_trial(trial, *fit_args, **fit_kwargs)
  File "/home/tmasters/anaconda3/envs/pacmann/lib/python3.10/site-packages/keras_tuner/src/engine/base_tuner.py", line 239, in _run_and_update_trial
    results = self.run_trial(trial, *fit_args, **fit_kwargs)
  File "/home/tmasters/anaconda3/envs/pacmann/lib/python3.10/site-packages/keras_tuner/src/tuners/hyperband.py", line 427, in run_trial
    return super().run_trial(trial, *fit_args, **fit_kwargs)
  File "/home/tmasters/anaconda3/envs/pacmann/lib/python3.10/site-packages/keras_tuner/src/engine/tuner.py", line 314, in run_trial
    obj_value = self._build_and_fit_model(trial, *args, **copied_kwargs)
  File "/home/tmasters/anaconda3/envs/pacmann/lib/python3.10/site-packages/keras_tuner/src/engine/tuner.py", line 232, in _build_and_fit_model
    model = self._try_build(hp)
  File "/home/tmasters/anaconda3/envs/pacmann/lib/python3.10/site-packages/keras_tuner/src/engine/tuner.py", line 167, in _try_build
    raise errors.FatalTypeError(
keras_tuner.src.errors.FatalTypeError: Expected the model-building function, or HyperModel.build() to return a valid Keras Model instance. Received: <tf_keras.src.engine.sequential.Sequential object at 0x1497ce38ebc0> of type <class 'tf_keras.src.engine.sequential.Sequential'>.

I am confused as to why the Sequential model is not being considered a "valid Keras Model." I would appreciate any help!