Repository: tensorflow
Issue: [TFLite C++] Signature calculating CategoricalCrossentropy loss produces wrong result
Labels: stat:awaiting tensorflower,type:bug,comp:lite,type:performance,TF 2.13

<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Bug

### Have you reproduced the bug with TF nightly?

Yes

### Source

source

### Tensorflow Version

2.13

### Custom Code

Yes

### OS Platform and Distribution

Windows 10

### Mobile device

_No response_

### Python version

_No response_

### Bazel version

5.3.0

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current Behaviour?

I've created a simple model in Python (TF version 2.10) and converted it for tflite. The model has two signatures, one for inference and other for training. When I run those signatures in Python, everything works correctly, I get good inference result and good training loss. When I load the converted tflite model with the C++ TFLite API (built from source, from branch r2.13) and run those signatures: inference works as intended, training works as intended (the accuracy on the test set is steadily rising), but the reported loss is totally random. At first I thought that loss might be accumulated since it is rising to five digits, but that is not the case since it rises and falls in a random fashion. It seems like there is some bug in the ops used for CategoricalCrossentropy C++ TFLite implementation.

I've tried building tensorflow from r2.12 and r2.13 and I get the same behavior. I've tried r2.10 also but then I couldn't even run the signatures with C++ TFLite API, I was getting bunch of segmentation faults. I couldn't find anywhere the documentation on what ops for backward prop are available in C++ TFLite API, maybe some of those which are used in CategoricalCrossentropy loss calculation are not yet available, or there is a bug in their implementation.

### Standalone code to reproduce the issue

Here is a Python code I am using to create model with signatures:
```
IMG_SIZE = 28

class Model(tf.Module):
    def __init__(self):
        self.model = tf.keras.Sequential([
            tf.keras.layers.Flatten(input_shape=(IMG_SIZE, IMG_SIZE), name='flatten'),
            tf.keras.layers.Dense(
                units=10,
                kernel_initializer=tf.keras.initializers.RandomNormal(mean=0.0, stddev=0.05),
                bias_initializer=tf.keras.initializers.Ones(),
                name='dense'
            ),
        ])

        opt = tf.keras.optimizers.SGD(learning_rate=0.1)
        loss_fn = tf.keras.losses.CategoricalCrossentropy(from_logits=True)
        self.model.compile(optimizer=opt, loss=loss_fn, metrics=['accuracy'])

    # The `train` function takes a batch of input images and labels.
    @tf.function(input_signature=[
        tf.TensorSpec([32, IMG_SIZE, IMG_SIZE], tf.float32),
        tf.TensorSpec([32, 10], tf.float32),
    ])
    def train(self, x, y):
        with tf.GradientTape() as tape:
            prediction = self.model(x)
            loss = self.model.loss(y, prediction)
        gradients = tape.gradient(loss, self.model.trainable_variables)
        self.model.optimizer.apply_gradients(
            zip(gradients, self.model.trainable_variables))
        result = {"loss": loss}
        return result

    @tf.function(input_signature=[
        tf.TensorSpec([1, IMG_SIZE, IMG_SIZE], tf.float32),
    ])
    def infer(self, x):
        logits = self.model(x)
        probabilities = tf.nn.softmax(logits, axis=-1)
        return {
            "output": probabilities,
            "logits": logits
        }
```

And here is the C++ code I am using to run the tflite model:
```
std::unique_ptr<tflite::FlatBufferModel> model =
    tflite::FlatBufferModel::BuildFromFile(tflite_model_path);
if (model == nullptr)
{
    std::cout << "Failed to load model" << std::endl;
    return;
}

tflite::ops::builtin::BuiltinOpResolver resolver;
tflite::InterpreterBuilder builder(*model, resolver);
std::unique_ptr<tflite::Interpreter> interpreter;
builder(&interpreter);
if (interpreter == nullptr)
{
    std::cout << "Failed to create interpreter" << std::endl;
    return;
}

if (interpreter->AllocateTensors() != kTfLiteOk)
{
    std::cout << "Failed to alocate interpreter tensors" << std::endl;
    return;
}

tflite::SignatureRunner* train_runner = interpreter->GetSignatureRunner("train");

TfLiteTensor* input_data_tensor = train_runner->input_tensor(train_runner->input_names()[0]);
float* input_data = input_data_tensor->data.f;
TfLiteTensor* input_labels_tensor = train_runner->input_tensor(train_runner->input_names()[1]);
float* input_labels = input_labels_tensor->data.f;

// Here I fill in the input data and labels, code redacted for brevity.

if (train_runner->Invoke() != kTfLiteOk)
{
    std::cout << "Error invoking train interpreter signature" << std::endl;
    return;
}

const TfLiteTensor* output_tensor = train_runner->output_tensor(train_runner->output_names()[0]);
float* output = output_tensor->data.f;
std::cout << "Training finished with loss: " << output[0] << std::endl;
```

Please let me know if you need more details, or full source code.


### Relevant log output

Here are the losses from batch to batch, as you can see they are too high and pretty much random. I repeat: the model is training correctly which I can see because the accuracy on the test set is steadily rising, so these loss values do not make sense.
```
Training of batch 1 finished with loss: 172.813
Training of batch 2 finished with loss: 30406.2
Training of batch 3 finished with loss: 35372.7
Training of batch 4 finished with loss: 30955.9
Training of batch 5 finished with loss: 30645.5
Training of batch 6 finished with loss: 39069.4
Training of batch 7 finished with loss: 25181.5
Training of batch 8 finished with loss: 28106.7
Training of batch 9 finished with loss: 12969.1
Training of batch 10 finished with loss: 3079.69
Training of batch 11 finished with loss: 3693.12
Training of batch 12 finished with loss: 3314.77
Training of batch 13 finished with loss: 4591.12
Training of batch 14 finished with loss: 5880.76
Training of batch 15 finished with loss: 5654.75
Training of batch 16 finished with loss: 10133.1
Training of batch 17 finished with loss: 9301.94
Training of batch 18 finished with loss: 11654.5
Training of batch 19 finished with loss: 11827.8
Training of batch 20 finished with loss: 22028.1
Training of batch 21 finished with loss: 8553.58
```

</details>