Q:

In tensorflow/keras, why train_on_batch output loss is different when recomputed after training using predict_on_batch?

I created a model using keras and I trained it using train_on_batch. To check if the model does what it's supposed to do, I recomputed the loss before and after the training phase using the predict_on_batch method. But, as you'd guess when reading the title, I don't have the same output losses.
Here follows a basic code to illustrate my problem:
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
import tensorflow as tf
import numpy as np 

# Loss definition
def mse(y_true, y_pred):
    return tf.reduce_mean(tf.square(y_true-y_pred))

# Model definition
model = Sequential()
model.add(Dense(1))
model.compile('rmsprop',mse)

# Data creation
batch_size = 10
x = np.random.random_sample([batch_size,10])
y = np.random.random_sample(batch_size)

# Print loss before training
y_pred = model.predict_on_batch(x)
print("Before: " + str(mse(y,y_pred).numpy()))

# Print loss output from train_on_batch
print("Train output: " + str(model.train_on_batch(x,y)))

# Print loss after training
y_pred = model.predict_on_batch(x)
print("After: " + str(mse(y,y_pred).numpy()))

With this code I get the following outputs:
Before: 0.28556848
Train output: 0.29771945
After: 0.27345362

I'd suppose that the training loss and the loss computed after training should be the same. So I'd like to understand why not?

A:

This is how train_on_batch works, it calculates the loss, then updates the network, so we get the loss before the network was updated.
When we apply predict_on_batch, we get the prediction from the updated network.
Under the hood, train_on_batch does many more things like fixing your data types, standardizing your data, etc.
The closest sibling of train_on_batch would be test_on_batch. If you run test_on_batch you'll find the result is close to train_on_bacth but not same.
Here's an implementation of test_on_batch : https://github.com/tensorflow/tensorflow/blob/e5bf8de410005de06a7ff5393fafdf832ef1d4ad/tensorflow/python/keras/engine/training_v2_utils.py#L442
it internally calls _standardize_user_data to fix your data types, data shapes, etc.
Once, you fix your x and y with proper shapes and data types the result is very close except for some small difference delta due to numerical instability.
Here's a minimal example where the test_on_batch, train_on_batch and predict_on_batch seem to agree on the result numerically.
from tensorflow.keras.layers import *
from tensorflow.keras.models import Model, Sequential
from tensorflow.keras.optimizers import Adam
import tensorflow as tf
import numpy as np

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense
import tensorflow as tf
import numpy as np 

# Loss definition
def mse(y_true, y_pred):
    return tf.reduce_mean(tf.square(y_true-y_pred))

# Model definition
model = Sequential()
model.add(Dense(1, input_shape = (10,)))
model.compile(optimizer = 'adam', loss = mse, metrics = [mse])

# Data creation
batch_size = 10
x = np.random.random_sample([batch_size,10]).astype('float32').reshape(-1, 10)
y = np.random.random_sample(batch_size).astype('float32').reshape(-1,1)

print(x.shape)
print(y.shape)

model.summary()

# running 5 iterations to check
for _ in range(5):

  # Print loss before training
  y_pred = model.predict_on_batch(x)
  print("Before: " + str(mse(y,y_pred).numpy()))

  # Print loss output from train_on_batch
  print("Train output: " + str(model.train_on_batch(x,y)))

  print(model.test_on_batch(x, y))

  # Print loss after training
  y_pred = model.predict_on_batch(x)
  print("After: " + str(mse(y,y_pred).numpy()))

(10, 10)
(10, 1)
Model: "sequential_25"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
dense_27 (Dense)             (None, 1)                 11        
=================================================================
Total params: 11
Trainable params: 11
Non-trainable params: 0
_________________________________________________________________
Before: 0.30760005
Train output: [0.3076000511646271, 0.3076000511646271]
[0.3052913546562195, 0.3052913546562195]
After: 0.30529135
Before: 0.30529135
Train output: [0.3052913546562195, 0.3052913546562195]
[0.30304449796676636, 0.30304449796676636]
After: 0.3030445
Before: 0.3030445
Train output: [0.30304449796676636, 0.30304449796676636]
[0.3008604645729065, 0.3008604645729065]
After: 0.30086046
Before: 0.30086046
Train output: [0.3008604645729065, 0.3008604645729065]
[0.2987399995326996, 0.2987399995326996]
After: 0.29874
Before: 0.29874
Train output: [0.2987399995326996, 0.2987399995326996]
[0.2966836094856262, 0.2966836094856262]
After: 0.2966836

N.B: train_on_batch updates the weight of the neural network after calculating the loss, so obviously the loss from train_on_batch and test_on_batch or predict_on_batch won't be same exactly. The proper question could be why test_on_batch and predict_on_batch give different losses with your data.

