Split a tensoflow model into two parts
After several attemps I have the following piece of code. I created two new networks and transferred the weights from the initial network to them.

# first part of initial model
part1_model = tf.keras.models.Sequential([
  tf.keras.layers.Conv2D(32, (3,3), activation='relu', input_shape=(28, 28, 1)),
  tf.keras.layers.MaxPooling2D(4, 4),
  tf.keras.layers.Conv2D(32, (3,3), activation='relu'),
  tf.keras.layers.MaxPooling2D(2,2),
  tf.keras.layers.Flatten()
])
part1_model.layers[0].set_weights(model.layers[0].get_weights())
part1_model.layers[1].set_weights(model.layers[1].get_weights())
part1_model.layers[2].set_weights(model.layers[2].get_weights())
part1_model.layers[3].set_weights(model.layers[3].get_weights())
part1_model.layers[4].set_weights(model.layers[4].get_weights())
part1_model.summary() 

# second part of initial model
part2_model = tf.keras.models.Sequential([
    tf.keras.layers.Dense(128, activation='relu',input_shape=(1, 128)),
  tf.keras.layers.Dense(10, activation='softmax')
])
part2_model.summary()
part2_model.layers[0].set_weights(model.layers[5].get_weights())
part2_model.layers[1].set_weights(model.layers[6].get_weights())
# predictions holds prediction for test set from initial model
predictions = model.predict(test_images)

# predictions1 holds output of first part when test set is used as input 
predictions1 = part1_model.predict(test_images)

import numpy as np
# tmp is used to transform prediction1 in a format recognizable from part2
tmp=np.zeros((10000,1,128))
for i in range(0,10000):
    tmp[i,:,:]=predictions1[i,:]

# predictions2 holds the output of second part when the result of first part (predictions1) is used as input    
predictions2 = part2_model.predict(tmp)

# check that the result of initial model is the same with the result of the two parts
ok=0
for i in range(0,10000):
    if np.argmax(predictions[i])!=np.argmax(predictions2[i]):
        print(i," False")
    else:
        ok=ok+1
print(ok)
# also sample some cases and check probabilities
print("last entry")
print(predictions2[9999])
print(predictions[9999])
print("2nd entry")
print(predictions2[1])
print(predictions[1])