An error occurred while executing the following cell:
------------------
fixed_vars = [None, 1]

for fixed_var in fixed_vars:
    model = DensityNetwork([1, 50, 50], [20, 1], fixed_var)

    def negloglik(y, rv_y):
        return -rv_y.log_prob(y)

    model.compile(optimizer=tf.optimizers.Adam(learning_rate=0.01), loss=negloglik)
    history = model.fit(x, y, epochs=500, batch_size=100, verbose=False)
    if fixed_var is not None:
        # estimate MLE of sigma
        yhat_train = model(x)
        ypred = yhat_train.mean().numpy()[:, 0]
        # print(ypred.shape)
        residuals = y - ypred
        mse = np.mean((residuals**2))
        # mse = scipy.stats.trim_mean(residuals ** 2, proportiontocut=0.1)
        model.fixed_variance = mse
        # print(model.fixed_variance)
    yhat = model(x_tst)  # a Gaussian distribution object

    plt.plot(history.history["loss"], label="Train")
    plt.legend()
    plt.xlabel("Epoch")
    plt.ylabel("NLL")
    plt.show()

    plt.figure()
    plt.plot(x, y, "b.", label="observed")
    m = yhat.mean()
    s = yhat.stddev()
    plt.plot(x_tst, m, "r", linewidth=4, label="mean")
    plt.plot(x_tst, m + 2 * s, "g", linewidth=2, label=r"mean + 2 stddev")
    plt.plot(x_tst, m - 2 * s, "g", linewidth=2, label=r"mean - 2 stddev")
    # plt.ylim([-10,10])
    plt.show()
------------------

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
/tmp/ipykernel_5298/2370705851.py in <module>
      8 
      9     model.compile(optimizer=tf.optimizers.Adam(learning_rate=0.01), loss=negloglik)
---> 10     history = model.fit(x, y, epochs=500, batch_size=100, verbose=False)
     11     if fixed_var is not None:
     12         # estimate MLE of sigma

~/miniconda3/envs/py37/lib/python3.7/site-packages/keras/utils/traceback_utils.py in error_handler(*args, **kwargs)
     68             # To get the full stack trace, call:
     69             # `tf.debugging.disable_traceback_filtering()`
---> 70             raise e.with_traceback(filtered_tb) from None
     71         finally:
     72             del filtered_tb

~/miniconda3/envs/py37/lib/python3.7/site-packages/keras/engine/training.py in tf__train_function(iterator)
     13                 try:
     14                     do_return = True
---> 15                     retval_ = ag__.converted_call(ag__.ld(step_function), (ag__.ld(self), ag__.ld(iterator)), None, fscope)
     16                 except:
     17                     do_return = False

AttributeError: in user code:

    File "/github/home/miniconda3/envs/py37/lib/python3.7/site-packages/keras/engine/training.py", line 1249, in train_function  *
        return step_function(self, iterator)
    File "/github/home/miniconda3/envs/py37/lib/python3.7/site-packages/keras/engine/training.py", line 1233, in step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "/github/home/miniconda3/envs/py37/lib/python3.7/site-packages/keras/engine/training.py", line 1222, in run_step  **
        outputs = model.train_step(data)
    File "/github/home/miniconda3/envs/py37/lib/python3.7/site-packages/keras/engine/training.py", line 1024, in train_step
        loss = self.compute_loss(x, y, y_pred, sample_weight)
    File "/github/home/miniconda3/envs/py37/lib/python3.7/site-packages/keras/engine/training.py", line 1083, in compute_loss
        y, y_pred, sample_weight, regularization_losses=self.losses
    File "/github/home/miniconda3/envs/py37/lib/python3.7/site-packages/keras/engine/compile_utils.py", line 263, in __call__
        y_t, y_p, sw = match_dtype_and_rank(y_t, y_p, sw)
    File "/github/home/miniconda3/envs/py37/lib/python3.7/site-packages/keras/engine/compile_utils.py", line 831, in match_dtype_and_rank
        if y_t.shape.rank == 1 and y_p.shape.rank == 2:

    AttributeError: 'Normal' object has no attribute 'shape'

AttributeError: in user code:

    File "/github/home/miniconda3/envs/py37/lib/python3.7/site-packages/keras/engine/training.py", line 1249, in train_function  *
        return step_function(self, iterator)
    File "/github/home/miniconda3/envs/py37/lib/python3.7/site-packages/keras/engine/training.py", line 1233, in step_function  **
        outputs = model.distribute_strategy.run(run_step, args=(data,))
    File "/github/home/miniconda3/envs/py37/lib/python3.7/site-packages/keras/engine/training.py", line 1222, in run_step  **
        outputs = model.train_step(data)
    File "/github/home/miniconda3/envs/py37/lib/python3.7/site-packages/keras/engine/training.py", line 1024, in train_step
        loss = self.compute_loss(x, y, y_pred, sample_weight)
    File "/github/home/miniconda3/envs/py37/lib/python3.7/site-packages/keras/engine/training.py", line 1083, in compute_loss
        y, y_pred, sample_weight, regularization_losses=self.losses
    File "/github/home/miniconda3/envs/py37/lib/python3.7/site-packages/keras/engine/compile_utils.py", line 263, in __call__
        y_t, y_p, sw = match_dtype_and_rank(y_t, y_p, sw)
    File "/github/home/miniconda3/envs/py37/lib/python3.7/site-packages/keras/engine/compile_utils.py", line 831, in match_dtype_and_rank
        if y_t.shape.rank == 1 and y_p.shape.rank == 2:

    AttributeError: 'Normal' object has no attribute 'shape'