probabilities = model(features, training=False)
row_sums = ops.convert_to_numpy(ops.sum(probabilities, axis=1))

assert model.input_shape == (None, 4)
assert model.output_shape == (None, 3)
assert probabilities.shape == (2, 3)
np.testing.assert_allclose(row_sums, np.ones(2), rtol=1e-5, atol=1e-5)

print(f"backend: {keras.config.backend()}")
print(f"model: {model.name}")
print(f"layers: {[layer.name for layer in model.layers]}")
print(f"parameters: {model.count_params()}")
print(f"input shape: {model.input_shape}")
print(f"output shape: {model.output_shape}")
print(f"prediction shape: {probabilities.shape}")
print(f"probability sums: {row_sums.round(6).tolist()}")