Repository: tensorflow
Issue: TypeError: true_fn and false_fn arguments to tf.cond must have the same number, type, and overall structure of return values
Labels: type:bug,2.17

### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

No

### Source

binary

### TensorFlow version

2.17.0

### Custom code

Yes

### OS platform and distribution

Linux CentOS 7.9

### Mobile device

_No response_

### Python version

3.12.4

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

We used tensorflow for 3 class classifications. The code worked well on existed environment: Python 3.8.3 + tensorflow 2.13.1.
But when we tried on new enviroment: Python 3.12.4 + tensorflow 2.17.0, several errors blocked.

### Standalone code to reproduce the issue


Here are the simple code to reproduce the issue.
```
import tensorflow as tf
from tensorflow.keras.layers import Dense
import numpy as np


def build_model():
    metrics = ['accuracy', tf.keras.metrics.AUC(multi_label=True)]
    loss = tf.keras.losses.sparse_categorical_crossentropy
    num_class = 3
    inputs = tf.keras.layers.Input(shape=(40, 36))
    lstm_out = tf.keras.layers.LSTM(32, return_sequences=False)(inputs)
    output = Dense(num_class, activation='softmax')(lstm_out)
    model = tf.keras.Model(inputs, output)
    model.compile(loss=loss, optimizer=tf.keras.optimizers.Adam(), metrics=metrics)
    model.summary()
    return model

def data_gen():
    while True:
        yield np.random.rand(128, 40, 36), np.random.randint(0, 3, (128, 1))

model = build_model()
cw = {0: 0.3, 1:2.5, 2:3.2}
model.fit(data_gen(), epochs=2, steps_per_epoch=10, class_weight=cw)
```

The first error we met:
```
Traceback (most recent call last):
  File "/data/release/infinity_stock2/tf_model.py", line 25, in <module>
    model.fit(data_gen(), epochs=2, steps_per_epoch=10, class_weight=cw)
  File "/root/.virtualenvs/new_stock/lib/python3.12/site-packages/keras/src/utils/traceback_utils.py", line 122, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/root/.virtualenvs/new_stock/lib/python3.12/site-packages/keras/src/trainers/data_adapters/__init__.py", line 108, in get_data_adapter
    raise ValueError(
ValueError: Argument `class_weight` is not supported for Python generator inputs. Received: class_weight={0: 0.3, 1: 2.5, 2: 3.2}
```

To solve the issue, we used tf dataset to wrap the generator. 
```
import tensorflow as tf
from tensorflow.keras.layers import Dense
import numpy as np


def build_model():
    metrics = ['accuracy', tf.keras.metrics.AUC(multi_label=True)]
    loss = tf.keras.losses.sparse_categorical_crossentropy
    num_class = 3
    inputs = tf.keras.layers.Input(shape=(40, 36))
    lstm_out = tf.keras.layers.LSTM(32, return_sequences=False)(inputs)
    output = Dense(num_class, activation='softmax')(lstm_out)
    model = tf.keras.Model(inputs, output)
    model.compile(loss=loss, optimizer=tf.keras.optimizers.Adam(), metrics=metrics)
    model.summary()
    return model

def data_gen():
    while True:
        yield np.random.rand(128, 40, 36), np.random.randint(0, 3, (128, 1))

model = build_model()
cw = {0: 0.3, 1:2.5, 2:3.2}
data = tf.data.Dataset.from_generator(data_gen, output_types=(tf.float32, tf.float32),  output_shapes=((None, 40, 36), (None, None)))
model.fit(data, epochs=2, steps_per_epoch=10, class_weight=cw)
```
And another error happened:
```
Traceback (most recent call last):
  File "/data/release/infinity_stock2/tf_model.py", line 26, in <module>
    model.fit(data, epochs=2, steps_per_epoch=10, class_weight=cw)
  File "/root/.virtualenvs/new_stock/lib/python3.12/site-packages/keras/src/utils/traceback_utils.py", line 122, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/root/.virtualenvs/new_stock/lib/python3.12/site-packages/tensorflow/python/ops/cond_v2.py", line 876, in error
    raise TypeError(
TypeError: true_fn and false_fn arguments to tf.cond must have the same number, type, and overall structure of return values.

true_fn output: Tensor("cond/Identity:0", shape=(None,), dtype=int64)
false_fn output: Tensor("cond/Identity:0", shape=(None,), dtype=int32)

Error details:
Tensor("cond/Identity:0", shape=(None,), dtype=int64) and Tensor("cond/Identity:0", shape=(None,), dtype=int32) have different types
```
It looks like class weight might not work well.

### Relevant log output

_No response_