Bypass keras dimension check for custom loss - keras

Question

I am trying to implement my custom loss function but I am having difficulties making it work with Keras.
I am dealing with n-dimensional data, and I would like for each input sample to compute a loss based on a another n-dimensional vector, we'll call it P (it encodes the reliability of each corresponding input measures, to put it simply).
My model is an autoencoder with n inputs (X) and n outputs (Y). I started implementing a solution in which I would overload my Y matrix by appending it P : new_Y = concat([Y, P]), and then in my custom loss split this matrix back and compute my loss. Problem is, tensorflow seems to perform a shape check of the output of my model with my new_Y, before even looking at my loss function.
Code to reproduce the error :
import numpy as np
from tensorflow import keras
import keras.backend as K
from keras import Model
from keras.layers import Input, Dense
from keras.optimizers import Adam
from keras.losses import MeanSquaredError as TFMSE
tf_mse = TFMSE()
input_dim = 2
batch_size = 16
X = np.random.randn(input_dim * batch_size).reshape((batch_size, input_dim, 1))
P = np.random.randn(input_dim * batch_size).reshape((batch_size, input_dim, 1))
Y = X.copy()
new_Y = np.hstack([Y, P])
_input = Input((input_dim, 1), name='input')
elayer = Dense(input_dim, kernel_initializer='normal', activation='relu')(_input)
bottleneck = Dense(1, kernel_initializer='normal', activation='linear')(elayer)
dlayer = Dense(input_dim, kernel_initializer='normal', activation='relu')(bottleneck)
model = Model(inputs=_input, outputs=dlayer)
def adjusted_mse(y_true, y_pred):
# I am first trying to simply apply classic MSE on Y here :
# y_true = [Y, P], so y_true[:, :input_dim] = Y
return tf_mse(y_true[:, :input_dim], y_pred)
model.compile(
loss=adjusted_mse,
optimizer=Adam(),
metrics=['mse']
)
model.fit(X, new_Y, epochs=8, batch_size=4, validation_split=.1, verbose=True)
Error raised :
Traceback (most recent call last):
File "<string>", line 1, in <module>
File "***/lib/python3.10/site-packages/keras/utils/traceback_utils.py", line 70, in error_handler
raise e.with_traceback(filtered_tb) from None
File "/tmp/__autograph_generated_filew2svrbng.py", line 15, in tf__train_function
retval_ = ag__.converted_call(ag__.ld(step_function), (ag__.ld(self), ag__.ld(iterator)), None, fscope)
ValueError: in user code:
...
ValueError: Dimensions must be equal, but are 2 and 4 for '{{node SquaredDifference}} = SquaredDifference[T=DT_FLOAT](model/dense_2/Relu, IteratorGetNext:1)' with input shapes: [?,2,2], [?,4,1].
Anyone have any idea how I could at least bypass this verification step ? Or is there a simpler way to do that ? I was thinking about overloading X the same way and passing the P part through an identity layer then concatenate it back, but that is twisted..
Thank you