Repository: tensorflow
Issue: Custom Loss Function can't access multiple y_true inputs. 
Labels: stat:awaiting tensorflower,type:support,comp:dist-strat,TF 2.15

### Issue type

Support

### Have you reproduced the bug with TensorFlow Nightly?

Yes

### Source

source

### TensorFlow version

2.15.0

### Custom code

Yes

### OS platform and distribution

Ubuntu 22.04 LTS

### Mobile device

Google Colab

### Python version

3.10.12

### Bazel version

_No response_

### GCC/compiler version

11.4.0

### CUDA/cuDNN version

12.1

### GPU model and memory

_No response_

### Current behavior?

I am using .tfrecord files as DataInput, I made a custom function that gives 4 inputs in X and 2 inputs in y. 
I am unable to access the 2nd input in y in a custom loss function. 

Data Input code

```python

seq_map = {"A":1,"C":2,"G":3,"U":4,"START":4,"END":5,"EMPTY":0}
bracket_map = {"(":1,")":2,"[":3,"]":4,"{":5,"}":6,"<":7,">":8,".":9,"START":10,"END":11,"EMPTY":0}

def convert_and_pad(ex,Lmax:int=206,shift=True,sn_filter:bool=True):

    l = tf.shape(ex["seq"],tf.int32)[0]

    if not shift:
        shift = 0
    else:
        shift = tf.random.uniform(shape=[1],minval=0,maxval=Lmax-l+1,dtype=tf.int32)[0]

    # Sequence Processing and Mask Processing
    seq = ex["seq"] + 1
    mask = tf.math.greater(tf.pad(seq,[[shift+1,Lmax-l-shift+1]]),seq_map["EMPTY"])
    seq = tf.pad(seq,[[1,0]],constant_values=seq_map["START"])                               # seq_map["START"]
    seq = tf.pad(seq,[[0,1]],constant_values=seq_map["END"])                                 # seq_map["END"]
    seq = tf.pad(seq,[[shift,Lmax-l-shift]])                                                 # seq_map["EMPTY"]


    # Bracket Processing
    brac = ex["bracket_seq"] + 1
    brac = tf.pad(brac,[[1,0]],constant_values=bracket_map["START"])                            # bracket_map["START"]
    brac = tf.pad(brac,[[0,1]],constant_values=bracket_map["END"])                              # bracket_map["END"]
    brac = tf.pad(brac,[[shift,Lmax-l-shift]],constant_values=bracket_map["EMPTY"])             # bracket_map["EMPTY"]

    # Reactivity Processing
    reac = tf.stack([ex["reactivity_2a3"][:l],ex["reactivity_dms"][:l]],axis=-1)
    reac = tf.pad(reac,[[shift+1,Lmax+1-l-shift],[0,0]],constant_values=np.nan)

    # SN_filter
    sn = (ex["sn_filter_2a3"] > 0) & (ex["sn_filter_dms"] > 0)

    # BPPMatrix
    bppm = ex["bpp_matrix"][:l,:l]
    bppm = tf.pad(bppm,[[shift+1,Lmax+1-l-shift],[shift+1,Lmax+1-l-shift]])

    seq.set_shape([208])
    brac.set_shape([208])

    if sn_filter:
        return (seq,brac,sn,bppm),(reac,mask)
    return (seq,brac,bppm),(reac,mask)


BATCH_SIZE = strategy.num_replicas_in_sync * 16 if (device == "TPU") else 32


def shape_set(X,y):
    (seq,brac,sn,bppm),(reac,mask) = X,y
    seq.set_shape([BATCH_SIZE,208])
    brac.set_shape([BATCH_SIZE,208])
    sn.set_shape([BATCH_SIZE])
    bppm.set_shape([BATCH_SIZE,208,208])
    reac.set_shape([BATCH_SIZE,208,2])
    mask.set_shape([BATCH_SIZE,208])
    return (seq,brac,sn,bppm),(reac,mask)


def create_train_ds(dataset,batch_size:int=BATCH_SIZE):
    dataset = dataset.map(convert_and_pad,num_parallel_calls=tf.data.AUTOTUNE)
    # dataset = dataset.repeat()
    dataset = dataset.shuffle(20000)
    dataset = dataset.batch(batch_size)
    dataset = dataset.map(lambda X,y:shape_set(X,y),num_parallel_calls=tf.data.AUTOTUNE)
    return dataset.prefetch(tf.data.AUTOTUNE)


def create_val_ds(dataset,batch_size:int=BATCH_SIZE):
    dataset = dataset.map(convert_and_pad,num_parallel_calls=tf.data.AUTOTUNE)
    dataset = dataset.batch(batch_size)
    dataset = dataset.map(lambda X,y:shape_set(X,y),num_parallel_calls=tf.data.AUTOTUNE)
    dataset = dataset.cache()
    return dataset.prefetch(tf.data.AUTOTUNE)

def create_test_ds(dataset,batch_size:int=BATCH_SIZE):
    dataset = dataset.map(convert_and_pad,num_parallel_calls=tf.data.AUTOTUNE)
    dataset = dataset.batch(batch_size)
    dataset = dataset.map(lambda X,y:shape_set(X,y),num_parallel_calls=tf.data.AUTOTUNE)
    return dataset.prefetch(tf.data.AUTOTUNE)


train_ds = create_train_ds(train_modified_ds,batch_size=BATCH_SIZE)
valid_ds = create_val_ds(val_modified_ds,batch_size=BATCH_SIZE)
test_ds = create_test_ds(test_modified_ds,batch_size=BATCH_SIZE)


```

```python

class CustomLoss(keras.losses.Loss):

    def __init__(self, **kwargs):
        super(CustomLoss, self).__init__(**kwargs)

    def __call__(self, y_true,y_pred,sample_weight=None):
        loss_y_true,loss_y_mask = y_true
        print(loss_y_true.shape,loss_y_mask.shape,y_pred.shape)
        loss_y_mask = tf.cast(loss_y_mask,tf.bool)
        loss_y_true = tf.clip_by_value(loss_y_true[loss_y_mask],clip_value_min=0,clip_value_max=1)
        loss_y_pred = tf.clip_by_value(y_pred[loss_y_mask],clip_value_min=0,clip_value_max=1)
        mae_loss = tf.math.reduce_mean(tf.math.abs(tf.math.subtract(loss_y_true,loss_y_pred)),axis=-1)
        return tf.math.reduce_mean(mae_loss[~tf.math.is_nan(mae_loss)])

```

Tried looking what was actually being accessed in custom loss

apparently only reactivity (reac) is being accessed as print(y_true.shape) gave me output of (32,208,2) so only true reac values are being accessed. I want to also have access to masked input values. 

### Standalone code to reproduce the issue

```shell
strategy = tf.distribute.OneDeviceStrategy("CPU")

def convert_and_pad(ex,Lmax:int=206,shift=True,sn_filter:bool=True):

    l = tf.shape(ex["seq"],tf.int32)[0]

    if not shift:
        shift = 0
    else:
        shift = tf.random.uniform(shape=[1],minval=0,maxval=Lmax-l+1,dtype=tf.int32)[0]

    # Sequence Processing and Mask Processing
    seq = ex["seq"] + 1
    mask = tf.math.greater(tf.pad(seq,[[shift+1,Lmax-l-shift+1]]),seq_map["EMPTY"])
    seq = tf.pad(seq,[[1,0]],constant_values=seq_map["START"])                               # seq_map["START"]
    seq = tf.pad(seq,[[0,1]],constant_values=seq_map["END"])                                 # seq_map["END"]
    seq = tf.pad(seq,[[shift,Lmax-l-shift]])                                                 # seq_map["EMPTY"]


    # Bracket Processing
    brac = ex["bracket_seq"] + 1
    brac = tf.pad(brac,[[1,0]],constant_values=bracket_map["START"])                            # bracket_map["START"]
    brac = tf.pad(brac,[[0,1]],constant_values=bracket_map["END"])                              # bracket_map["END"]
    brac = tf.pad(brac,[[shift,Lmax-l-shift]],constant_values=bracket_map["EMPTY"])             # bracket_map["EMPTY"]

    # Reactivity Processing
    reac = tf.stack([ex["reactivity_2a3"][:l],ex["reactivity_dms"][:l]],axis=-1)
    reac = tf.pad(reac,[[shift+1,Lmax+1-l-shift],[0,0]],constant_values=np.nan)

    # SN_filter
    sn = (ex["sn_filter_2a3"] > 0) & (ex["sn_filter_dms"] > 0)

    # BPPMatrix
    bppm = ex["bpp_matrix"][:l,:l]
    bppm = tf.pad(bppm,[[shift+1,Lmax+1-l-shift],[shift+1,Lmax+1-l-shift]])

    seq.set_shape([208])
    brac.set_shape([208])

    if sn_filter:
        return (seq,brac,sn,bppm),(reac,mask)
    return (seq,brac,bppm),(reac,mask)


BATCH_SIZE = strategy.num_replicas_in_sync * 16 if (device == "TPU") else 32


def shape_set(X,y):
    (seq,brac,sn,bppm),(reac,mask) = X,y
    seq.set_shape([BATCH_SIZE,208])
    brac.set_shape([BATCH_SIZE,208])
    sn.set_shape([BATCH_SIZE])
    bppm.set_shape([BATCH_SIZE,208,208])
    reac.set_shape([BATCH_SIZE,208,2])
    mask.set_shape([BATCH_SIZE,208])
    return (seq,brac,sn,bppm),(reac,mask)


def create_train_ds(dataset,batch_size:int=BATCH_SIZE):
    dataset = dataset.map(convert_and_pad,num_parallel_calls=tf.data.AUTOTUNE)
    # dataset = dataset.repeat()
    dataset = dataset.shuffle(20000)
    dataset = dataset.batch(batch_size)
    dataset = dataset.map(lambda X,y:shape_set(X,y),num_parallel_calls=tf.data.AUTOTUNE)
    return dataset.prefetch(tf.data.AUTOTUNE)


def create_val_ds(dataset,batch_size:int=BATCH_SIZE):
    dataset = dataset.map(convert_and_pad,num_parallel_calls=tf.data.AUTOTUNE)
    dataset = dataset.batch(batch_size)
    dataset = dataset.map(lambda X,y:shape_set(X,y),num_parallel_calls=tf.data.AUTOTUNE)
    dataset = dataset.cache()
    return dataset.prefetch(tf.data.AUTOTUNE)

def create_test_ds(dataset,batch_size:int=BATCH_SIZE):
    dataset = dataset.map(convert_and_pad,num_parallel_calls=tf.data.AUTOTUNE)
    dataset = dataset.batch(batch_size)
    dataset = dataset.map(lambda X,y:shape_set(X,y),num_parallel_calls=tf.data.AUTOTUNE)
    return dataset.prefetch(tf.data.AUTOTUNE)


train_ds = create_train_ds(train_modified_ds,batch_size=BATCH_SIZE)
valid_ds = create_val_ds(val_modified_ds,batch_size=BATCH_SIZE)
test_ds = create_test_ds(test_modified_ds,batch_size=BATCH_SIZE)

X,y = train_ds.take(1).get_single_element()
seq_input = X[0]
bracket_input = X[1]
sn_filter = X[2]
bpp_matrix = X[3]
reactivity = y[0]
mask = y[1]

print("sequence input shape: ",seq_input.shape)
print("bracket sequence input shape: ",bracket_input.shape)
print("sn filter input shape: ",sn_filter.shape)
print("Bpp matrix input shape: ",bpp_matrix.shape)
print("Reactivity input shape: ",reactivity.shape)
print("Mask input shape: ",mask.shape)

class CustomLoss(keras.losses.Loss):

    def __init__(self, **kwargs):
        super(CustomLoss, self).__init__(**kwargs)

    def __call__(self, y_true,y_pred,sample_weight=None):
        loss_y_true,loss_y_mask = y_true
        print(loss_y_true.shape,loss_y_mask.shape,y_pred.shape)
        loss_y_mask = tf.cast(loss_y_mask,tf.bool)
        loss_y_true = tf.clip_by_value(loss_y_true[loss_y_mask],clip_value_min=0,clip_value_max=1)
        loss_y_pred = tf.clip_by_value(y_pred[loss_y_mask],clip_value_min=0,clip_value_max=1)
        mae_loss = tf.math.reduce_mean(tf.math.abs(tf.math.subtract(loss_y_true,loss_y_pred)),axis=-1)
        return tf.math.reduce_mean(mae_loss[~tf.math.is_nan(mae_loss)])

## Testing function
cust_loss = CustomLoss()
y_pred = tf.random.uniform(shape=y[0].shape,minval=-1,maxval=2)
print("custom loss: ",cust_loss(y,y_pred).numpy())

with strategy.scope():
    model = Transformer()
    cust_loss = CustomLoss()
    cust_metric_2a3 = CustomMetric2A3()
    cust_metric_dms = CustomMetricDMS()
    model.compile(
        optimizer=keras.optimizers.Adam(1e-5),
        loss=cust_loss,
        metrics=[cust_metric_2a3,cust_metric_dms]
)
model.fit(train_ds,epochs=1)
```


### Relevant log output

```shell
sequence input shape:  (32, 208)
bracket sequence input shape:  (32, 208)
sn filter input shape:  (32,)
Bpp matrix input shape:  (32, 208, 208)
Reactivity input shape:  (32, 208, 2)
Mask input shape:  (32, 208)

custom loss: 0.47430646
---------------------------------------------------------------------------
OperatorNotAllowedInGraphError            Traceback (most recent call last)
<ipython-input-28-669a0a768bc0> in <cell line: 11>()
      9         metrics=[cust_metric_2a3,cust_metric_dms]
     10 )
---> 11 model.fit(train_ds,epochs=1)

1 frames
<ipython-input-27-bd2035bac2b8> in __call__(self, y_true, y_pred, sample_weight)
     50 
     51     def __call__(self, y_true,y_pred,sample_weight=None):
---> 52         loss_y_true,loss_y_mask = y_true
     53         print(loss_y_true.shape,loss_y_mask.shape,y_pred.shape)
     54         loss_y_mask = tf.cast(loss_y_mask,tf.bool)

OperatorNotAllowedInGraphError: Iterating over a symbolic `tf.Tensor` is not allowed. You can attempt the following resolutions to the problem: If you are running in Graph mode, use Eager execution mode or decorate this function with @tf.function. If you are using AutoGraph, you can try decorating this function with @tf.function. If that does not work, then you may be using an unsupported feature or your source code may not be visible to AutoGraph. See https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/g3doc/reference/limitations.md#access-to-source-code for more information.
```