mse loss function not compatible with regularization loss (add_loss) on hidden layer output

I would like to code in tf.Keras a Neural Network with a couple of loss functions. One is a standard mse (mean squared error) with a factor loading, while the other is basically a regularization term on the output of a hidden layer. This second loss is added through self.add_loss() in a user-defined class inheriting from tf.keras.layers.Layer. I have a couple of questions (the first is more important though).

1) The error I get when trying to combine the two losses together is the following:

ValueError: Shapes must be equal rank, but are 0 and 1
        From merging shape 0 with other shapes. for '{{node AddN}} = AddN[N=2, T=DT_FLOAT](loss/weighted_loss/value, model/new_layer/mul_1)' with input shapes: [], [100].

So it comes from the fact that the tensors which should add up to make one unique loss value have different shapes (and ranks). Still, when I try to print the losses during the training, I clearly see that the vectors returned as losses have shape batch_size and rank 1. Could it be that when the 2 losses are summed I have to provide them (or at least the loss of add_loss) as scalar? I know the mse is usually returned as a vector where each entry is the mse from one sample in the batch, hence having batch_size as shape. I think I tried to do the same with the "regularization" loss. Do you have an explanation for this behavio(u)r?

The sample code which gives me error is the following:

import numpy as np
import tensorflow as tf
from tensorflow.keras import backend as K
from tensorflow.keras.models import Model
from tensorflow.keras.layers import Dense, Input


def rate_mse(rate=1e5):
    @tf.function # also needed for printing
    def loss(y_true, y_pred):
        tmp = rate*K.mean(K.square(y_pred - y_true), axis=-1)
#        tf.print('shape %s and rank %s output in mse'%(K.shape(tmp), tf.rank(tmp)))
        tf.print('shape and rank output in mse',[K.shape(tmp), tf.rank(tmp)])
        tf.print('mse loss:',tmp) # print when I put tf.function
        return tmp
    return loss

class newLayer(tf.keras.layers.Layer):
    def __init__(self, rate=5e-2, **kwargs):
        super(newLayer, self).__init__(**kwargs)
        self.rate = rate
        
#    @tf.function # to be commented for NN training
    def call(self, inputs):
        tmp = self.rate*K.mean(inputs*inputs, axis=-1)
        tf.print('shape and rank output in regularizer',[K.shape(tmp), tf.rank(tmp)])
        tf.print('regularizer loss:',tmp)
        self.add_loss(tmp, inputs=True)
        return inputs



tot_n = 10000
xx = np.random.rand(tot_n,1)
yy = np.pi*xx

train_size = int(0.9*tot_n)
xx_train = xx[:train_size]; xx_val = xx[train_size:]
yy_train = yy[:train_size]; yy_val = yy[train_size:]

reg_layer = newLayer()

input_layer = Input(shape=(1,))                                      # input
hidden = Dense(20, activation='relu', input_shape=(2,))(input_layer) # hidden layer
hidden = reg_layer(hidden)
output_layer = Dense(1, activation='linear')(hidden) 

model = Model(inputs=[input_layer], outputs=[output_layer])
model.compile(optimizer='Adam', loss=rate_mse(), experimental_run_tf_function=False)
#model.compile(optimizer='Adam', loss=None, experimental_run_tf_function=False)
model.fit(xx_train, yy_train, epochs=100, batch_size = 100, 
          validation_data=(xx_val,yy_val), verbose=1)

#new_xx = np.random.rand(10,1); new_yy = np.pi*new_xx
#model.evaluate(new_xx,new_yy)


print(model.predict(np.array([[1]])))

2) I would also have a secondary question related to this code. I noticed that printing with tf.print inside the function rate_mse only works with tf.function. Similarly, the call method of newLayer is only taken into consideration if the same decorator is commented during training. Can someone explain why this is the case or reference me to a possible solution?

Thanks in advance to whoever can provide me help. I am currently using Tensorflow 2.2.0 and keras version is 2.3.0-tf.