title: Fidelity custom loss function in tensorflow - input shape problem
body: <p>I am having trouble with inputs to a custom loss function as part of a multi-objective elementary neural network (GRU so far). </p>
  
  <p>I have a multi-objective function that takes as input a flattened (,16) shape (4x4) matrix ("input matrix") and generates a sequence of 10 4x4 matrices (represented via a flattened output shape (,160)) as a first objective using 'mse' loss. This part of the model, which takes <code>data.lab</code> as training data and <code>inputreshape</code> as label data works fine. I need to, as a second objective, however convert this (,160) tensor to a (10,4,4) tensor (reshaping to 10 4x4 matrices), then matmul them together to obtain a (4,4) matrix "output matrix" (done via the productlayer custom layer). This output matrix then needs to be input into a function to calculate the 'fidelity' via the custom loss function 'fidelity2' by comparing with the original input from <code>data.lab</code>. I did this by inputting <code>data.lab</code> a second time, but now as a label. However (1) model.summary() seems to show the output of productlayer is (4,4) when I think it should be (None,4,4) and (2) I receive the error:</p>
  
  <p><code>InvalidArgumentError: In[0] is not a matrix. Instead it has shape [10,4,4]
       [[{{node loss_7/fidout_loss/ArithmeticOptimizer/FoldTransposeIntoMatMul_matmul}}]]</code></p>
  
  <p>The y_true in the custom loss function seems to be taking the entire label set from model.fit (see code below), rather than batching, so this means it dumps the (10,4,4) label set in but is trying to compare against the (4,4) y_pred. Unsure how to fix this problem. Code is below. Ultimately I need to be able to calculate fidelity - actually infidelity eventually - in this way for arbitrary batch sizes (I set batchsize in model.fit to 1 just to try to get it working). Suggestions appreciated.</p>
  
  <pre><code>def fidelity2(y_true, y_pred):
      y_truetp = tf.transpose(y_true)
      t1 = (y_truetp @ y_pred)
      tr = tf.trace(t1)
      mxdim = tf.cast(tf.shape(y_pred)[0], tf.float32)
      fidelity = (tf.abs((tr)** 2) / mxdim ** 2)
      return fidelity    
  
  
  
  x = layers.Input(shape=(data.realdim,data.realdim), name='input1', batch_size=None)
  x1 = layers.GRU(data.Uj_dim, return_sequences=True)(x)
  x1 = layers.Dropout(rate=0.2)(x1)
  x1 = layers.GRU(data.Uj_dim, return_sequences=True)(x1)
  x1 = layers.Dropout(rate=0.2)(x1)
  x1 = layers.GRU(data.Uj_dim, return_sequences=True)(x1)
  x1 = layers.Flatten()(x1)
  y = layers.Dense(160, activation='relu', name='output_data')(x1)
  xreshape = layers.Reshape((4,4))(x)
  y2 = productlayer(trainable = True, name="fidout")(x1)
  model = tf.keras.models.Model(inputs=x, outputs=[y,y2])
  
  
  
  batchsz = 1
  #===============
  model.compile(optimizer='adam', loss=['mse',fidelity2], metrics=['mse',fidelity2])
  model.summary()
  
  model.fit(data.lab, [inputreshape,data.lab], epochs=2, batch_size=batchsz, validation_split = 0, shuffle=False, steps_per_epoch=1)
  
  #==============```
  
  Model: "model_13"
  __________________________________________________________________________________________________
  Layer (type)                    Output Shape         Param #     Connected to                     
  ==================================================================================================
  input1 (InputLayer)             [(None, 4, 4)]       0                                            
  __________________________________________________________________________________________________
  gru_39 (GRU)                    (None, 4, 40)        5400        input1[0][0]                     
  __________________________________________________________________________________________________
  dropout_26 (Dropout)            (None, 4, 40)        0           gru_39[0][0]                     
  __________________________________________________________________________________________________
  gru_40 (GRU)                    (None, 4, 40)        9720        dropout_26[0][0]                 
  __________________________________________________________________________________________________
  dropout_27 (Dropout)            (None, 4, 40)        0           gru_40[0][0]                     
  __________________________________________________________________________________________________
  gru_41 (GRU)                    (None, 4, 40)        9720        dropout_27[0][0]                 
  __________________________________________________________________________________________________
  flatten_13 (Flatten)            (None, 160)          0           gru_41[0][0]                     
  __________________________________________________________________________________________________
  output_data (Dense)             (None, 160)          25760       flatten_13[0][0]                 
  __________________________________________________________________________________________________
  fidout (productlayer)           (4, 4)               0           flatten_13[0][0]                 
  ==================================================================================================
  Total params: 50,600
  Trainable params: 50,600
  Non-trainable params: 0
  __________________________________________________________________________________________________
  Epoch 1/2
  ---------------------------------------------------------------------------
  InvalidArgumentError...InvalidArgumentError: In[0] is not a matrix. Instead it has shape [10,4,4]
       [[{{node loss_7/fidout_loss/ArithmeticOptimizer/FoldTransposeIntoMatMul_matmul}}]]
  </code></pre>