Loss calculation fails with multiple outputs due to the LossWrapper

[markomitos]

When trying to execute forward_pass of a multi-output model with a loss the loss calculation fails because of a call to squeeze_or_expands_to_same_rank call inside the LossWrapper, as the parameters passed are of type tuple and list instead of a singular Tensor. The function fails on the call of shape on the passed parameters:

keras/keras/src/losses/loss.py

Line 109 in bce176f

def squeeze_or_expand_to_same_rank(x1, x2, expand_rank_1=True):

Loss' implementation initially calls to tf.convert_to_tenzor before the squeeze call while LossWrapper omits this step. As the wrapper is executed first, losses are no longer handle multiple outputs.

keras/keras/src/losses/losses.py

Line 1291 in bce176f

>>> loss = keras.losses.mean_squared_error(y_true, y_pred)

There are two ways to fix this problem that shouldn't pose too big of a problem:

• add convert_to_tensor first in LossWrapper
• or remove the squeeze_or_expands_to_same_rank call alltogether and rely on proper loss function implementation

[nicolaspi]

This is potentially fixed by #20358 as the loss wrapper now maps squeeze_or_expands_to_same_rank on the structure.

[markomitos]

@nicolaspi your PR did certanly help, although it did not fix the issue entirerly. Now the issue arrises when one of the parameters is a tuple and the other is a list as the check for same structure fails inside this function:

keras/keras/src/tree/optree_impl.py

Line 107 in 533ea63

raise ValueError(

And the error returned suggests that the structure itself is the same bar the difference in collection types:
ValueError: aandb don't have the same structure. Received: structure of a=PyTreeSpec((*, *, *), NoneIsLeaf), structure of b=PyTreeSpec([*, *, *], NoneIsLeaf)

[nicolaspi]

@markomitos Can you provide a minimal reproducible example?

[markomitos]

@nicolaspi The error arrises when running this test with the keras3 multiple output model. I have commented out all other test parameters for you conviniece:
https://github.com/markomitos/tensorflow-federated/blob/c1a094b8107f147febd6c27c8c126df4c3b6a7dc/tensorflow_federated/python/learning/models/serialization_test.py#L552

[google-ml-butler]

Are you satisfied with the resolution of your issue?
Yes
No

[nicolaspi]

@markomitos can you check whether your issue is fixed with the changes?

[markomitos]

@nicolaspi unfortunately no, the same error persists

[nicolaspi]

Sorry, I tried my best. Unfortunately, I can't help you further unless you provide a self contained minimal reproducible example that is not embedded in a third party library.
Feel free to reopen the issue.

[markomitos]

@nicolaspi Here is a minimal reproducible example that only uses keras that triggers the same error:

  def test_keras_loss_serialization(self):
    model = model_examples.build_multiple_outputs_keras3_model()
    model.compile(
      optimizer='adam',
      loss=keras.losses.MeanSquaredError()
    )
    x_batch = [
      tf.constant([[1.0], [2.0]]),
      tf.constant([[3.0], [4.0]])
    ]
    y_true = (2.0,
      6.0,
      10.0)

    y_pred = model.predict_on_batch(x_batch)
    loss_fn = keras.losses.MeanSquaredError()
    loss = loss_fn(y_true, y_pred)

def build_multiple_outputs_keras3_model():
  """Builds a test model with three outputs."""
  l = keras.layers
  a = l.Input((1,))
  b = l.Input((1,))

  output_a = l.Dense(1)(a)
  output_b = l.Dense(1)(b)
  output_c = l.Dense(1)(l.concatenate([l.Dense(1)(a), l.Dense(1)(b)]))

  return keras.Model(inputs=[a, b], outputs=[output_a, output_b, output_c])

I don't see the button for the reopening of the issue, perhaps I dont have the permissions to do it. Could you please reopen the issue?

[nicolaspi]

Hi,
The error occurs because there's a structure mismatch between y_pred and y_true when calling the loss function API. The comparison checks not only the data structure but also the container type (tuple vs list).
To resolve this, you have two options:

• Ensure both y_pred and y_true have matching structures and container types
• Use the Trainer API instead, which includes built-in logic to handle such structural mismatches automatically.
  e.g.:

import keras
import tensorflow as tf

def build_multiple_outputs_keras3_model():
  """Builds a test model with three outputs."""
  l = keras.layers
  a = l.Input((1,))
  b = l.Input((1,))

  output_a = l.Dense(1)(a)
  output_b = l.Dense(1)(b)
  output_c = l.Dense(1)(l.concatenate([l.Dense(1)(a), l.Dense(1)(b)]))

  return keras.Model(inputs=[a, b], outputs=[output_a, output_b, output_c])

model = build_multiple_outputs_keras3_model()
model.compile(optimizer='adam', loss=keras.losses.MeanSquaredError())
x_batch = [tf.constant([[1.0], [2.0]]), tf.constant([[3.0], [4.0]])]
y_true = (2.0, 6.0, 10.0)
loss = model.test_on_batch(x_batch, y_true)
print(loss)