Very cool! Is there any way we can integrate the additional required step to call loss.set_specs into sth like the compile step or so? Right now only Loss objects can be used with the new feature, when you pass a callable, the usual single tensor specs are assumed.

I am wondering if we really need to tinker with the y_true and y_pred arguments to model.compute_loss so much? Let's consider the possible scenarios:

1. Single loss, inputs are single tensors: we don't need to do anything to inputs, pass them right through to the single loss
2. Single loss, inputs are nested structures of tensors: again, we don't need to do anything to inputs, pass them right through to the single loss
3. multiple losses, inputs are single tensors: we need to broadcast the inputs to each loss, but leave their structures as they are
4. multiple losses, inputs are nested structures of tensors: this one is the only tricky part, let's dive in one more level:
   1. inputs are dicts: compare keys with losses dicts, if everything matches, pass the values through to each loss unaltered
   2. inputs are nested sequences: this is the only part that can be ambiguous (I think): if both y_true and y_pred are sequences of the same length as the losses it could mean either:
      • each input element corresponds to one loss, so zip iterate through losses, y_true and y_pred
      • each loss takes the same nested input shapes, so broadcast

So, 4.ii. is the only case we have to figure out how to make unambiguous.

One easy solution would be to always require the loss argument to model.compile as well as the inputs to model.compute_loss to be a sequence (or dict) of same length, even for a single loss. It would be the same as for the metrics argument to model.compile, so I think this would be fine.

What do you think? Did I miss any scenarios?

Hey @Darkdragon84

Thanks for your inputs.

Is there any way we can integrate the additional required step to call loss.set_specs into sth like the compile step or so?

It is quite difficult for me. The root cause is that we have no information about the structure of y_true and y_pred during build. I don't have an idea to automatically infer that.

One easy solution would be to always require the loss argument to model.compile as well as the inputs to model.compute_loss to be a sequence (or dict) of same length, even for a single loss.

Yeah, this will make CompileLoss much cleaner and we can avoid using specs. However, this change could break a lot of existing projects built with Keras. I don't think it is feasible.

Yeah, this will make CompileLoss much cleaner and we can avoid using specs. However, this change could break a lot of existing projects built with Keras. I don't think it is feasible.

Yeah I understand, especially since the latest major release was very recently... Would it be possible to enable arbitrarily nested tensor inputs only if the losses and inputs are given as lists (i.e. confer to the format that the metrics also use?). Perhaps we can also find an easier way to resolve the ambiguity without having to require list inputs

It is quite difficult for me. The root cause is that we have no information about the structure of y_true and y_pred during build. I don't have an idea to automatically infer that.

If we found a way to disambiguate 4.ii. we could defer the call to set_specs to when the losses are actually compiled (first call of compute_loss I think?). I'll think a bit about it.

Would it be possible to enable arbitrarily nested tensor inputs only if the losses and inputs are given as lists

This should be feasible. It's the easiest scenario. (I can create a shortcut for this)
On second thought, it isn't clear to me how we can verify that the inputs are valid even with a single loss

I tried to build this testing table and we should pass all of them with the algorithm

All possible valid scenario in tests

We should be careful that both y1 and y2 could be nested structure.

There are also some cases that should be considered invalid such as missed keys, but I haven't listed them above to simplify the table.

I don't think we should make changes to the base Loss API (i.e. let's restricts changes to the CompileLoss class). I also don't think we should make loss.set_specs() a public-facing API that users have to call if their losses take more than a single tensor input.

If we're going to enable nested structures for y_true and y_pred in losses, they should "just work". There should not be an additional config step that users have to take.

we could defer the call to set_specs to when the losses are actually compiled (first call of compute_loss I think?

This would be the right approach -- if we need to keep track of input specs, let's base it on the structure of the inputs received at the first call to the loss.

we could defer the call to set_specs to when the losses are actually compiled (first call of compute_loss I think?

This would be the right approach -- if we need to keep track of input specs, let's base it on the structure of the inputs received at the first call to the loss.

But how can we do this?
As far as I know, CompileLoss is the very first point where loss and the inputs (y_true, y_pred) meet, not to mention that the structures might already be nested.

(I'm going to separate the unrelated fix from this PR to concentrate the discussion here)

As far as I know, CompileLoss is the very first point where loss and the inputs (y_true, y_pred) meet, not to mention that the structures might already be nested.

So could the logic live entirely in CompileLoss and not have any user-facing footprint?

I have not thought about how to do it, but if the only way to make the feature work is to add a user facing set_specs step, then I think we're better off not supporting the feature at all by default (this would mean that users who need nested loss inputs will write their own compute_loss function).

But how can we do this?
As far as I know, CompileLoss is the very first point where loss and the inputs (y_true, y_pred) meet, not to mention that the structures might already be nested.

@james77777778
I guess at exactly that point we need to find a way to clearly infer from the inputs. If there are disambiguities we cannot resolve we probably need to retain the original behavior. Are you concerned about build?

So could the logic live entirely in CompileLoss and not have any user-facing footprint?

I think this is not feasible. We should retain the original behavior.

Are you concerned about build?

Hey @Darkdragon84. I'm unsure what do you mean by build?

I'm thinking of another way to support nested y_true and y_pred by utilizing dict. It should be feasible to check whether loss, y_true and y_pred are all dicts or not. If they are all dicts, we can directly feed the inputs to the corresponding loss, no matter what their structure is.
(Maybe it is possible to extend this approach a little more to support other structures as well.)

A simple usage like this:

import keras


def my_loss(y_true, y_pred):
    y1, y2 = y_true
    pred_sum = keras.ops.sum(y_pred)
    return keras.ops.abs(keras.ops.sum(y1) - pred_sum) + keras.ops.abs(
        keras.ops.sum(y2) - pred_sum
    )


x = keras.Input(shape=(3,), name="input_a")
output_a = keras.layers.Dense(1, name="output_a")(x)
output_b = keras.layers.Dense(1, name="output_b", activation="sigmoid")(x)
model = keras.Model(x, {"output": [output_a, output_b]})  # <- dict

model.compile(loss={"output": my_loss})  # <- dict

x = keras.random.uniform([8, 3])
y1 = keras.random.uniform([8, 1])
y2 = keras.random.randint((8, 1), 0, 2)
model.fit(x, {"output": [y1, y2]}, batch_size=2, epochs=1)  # <- dict

"""
Currently, this causes an error:
ValueError: In the dict argument `loss`, key 'output' does not correspond to any model output. Received:
loss={'output': <function my_loss at 0x73e7325fe2a0>}
"""