SPACE_TO_BATCH_ND output shape mismatch with TFLite

[jianyuzzz]

When testing the SPACE_TO_BATCH_ND operator, I noticed that TFLite and TFLite Micro produce different output shapes for the same model and input. There is no error from TFLite Micro, but the output shape does not match TFLite (and Keras).

Minimal Example:

import tensorflow as tf

# Minimal Keras model using SPACE_TO_BATCH_ND
inputs = tf.keras.Input(shape=(8, 8, 1))
block_shape = [2, 2]
paddings = [[0, 0], [0, 0]]
outputs = tf.keras.layers.Lambda(lambda x: tf.space_to_batch_nd(x, block_shape, paddings))(inputs)
model = tf.keras.Model(inputs, outputs)

# Convert to TFLite
converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()
with open("space_to_batch_nd.tflite", "wb") as f:
    f.write(tflite_model)

Inference Comparison:

import numpy as np
from ai_edge_litert.interpreter import Interpreter
from tflite_micro.python.tflite_micro import runtime

# Load TFLite model and prepare input
model_path = "space_to_batch_nd.tflite"
interpreter = Interpreter(model_path=model_path)
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
input_shape = input_details[0]["shape"]
input_data = np.random.rand(*input_shape).astype(np.float32)

# TFLite inference
interpreter.set_tensor(input_details[0]["index"], input_data)
interpreter.invoke()
tflite_output = interpreter.get_tensor(interpreter.get_output_details()[0]["index"])

# TFLM inference
arena_size = 1000000
tflm_interpreter = runtime.Interpreter.from_file(model_path, arena_size=arena_size)
tflm_interpreter.set_input(input_data, 0)
tflm_interpreter.invoke()
tflm_output = tflm_interpreter.get_output(0)

print("TFLite output shape:", tflite_output.shape)
print("TFLM output shape:", tflm_output.shape)

Observed Behavior:

• TFLite output shape: (4, 4, 4, 1)
• TFLite Micro output shape: (1, 4, 4, 1)

Environment:

• TensorFlow version: 2.18.1
• tflite-micro version: 0.dev20250812203306
• Platform: Linux

[ddavis-2015]

@jianyuzzz Try the following instead when converting the model:

import tensorflow as tf

# Minimal Keras model using SPACE_TO_BATCH_ND
inputs = tf.keras.Input(shape=(8, 8, 1))
block_shape = [2, 2]
paddings = [[0, 0], [0, 0]]
outputs = tf.keras.layers.Lambda(lambda x: tf.space_to_batch_nd(x, block_shape, paddings))(inputs)
model = tf.keras.Model(inputs, outputs)

# Get concrete function
model_func = tf.function(func=model)
cf = model_func.get_concrete_function(
          tf.TensorSpec(shape=(1,)+(8, 8, 1), dtype=tf.float32))

# Convert to TFLite
converter = tf.lite.TFLiteConverter.from_concrete_functions([cf],trackable_obj=model)
tflite_model = converter.convert()
tf.lite.experimental.Analyzer.analyze(model_content=tflite_model)

You will get this output, which shows the output shape to be fully computed (T#3):

=== TFLite ModelAnalyzer ===

Your TFLite model has '1' subgraph(s). In the subgraph description below,
T# represents the Tensor numbers. For example, in Subgraph#0, the SPACE_TO_BATCH_ND op takes
tensor #0 and tensor #1 and tensor #2 as input and produces tensor #3 as output.

Subgraph#0 main(T#0) -> [T#3]
  Op#0 SPACE_TO_BATCH_ND(T#0, T#1[2, 2], T#2[0, 0, 0, 0]) -> [T#3]

Tensors of Subgraph#0
  T#0(args_0) shape:[1, 8, 8, 1], type:FLOAT32
  T#1(functional_13_1/lambda_14_1/SpaceToBatchND/block_shape) shape:[2], type:INT32 RO 8 bytes, buffer: 2, data:[2, 2]
  T#2(functional_13_1/lambda_14_1/SpaceToBatchND/paddings) shape:[2, 2], type:INT32 RO 16 bytes, buffer: 3, data:[0, 0, 0, 0]
  T#3(Identity) shape:[4, 4, 4, 1], type:FLOAT32

[ddavis-2015]

[internal: @veblush try recomputing the output shape, again? This was done in #2335 which was reverted in #2339 in favor of #2356 which itself was reverted in #2390 ... 😭]

[jianyuzzz]

Sorry for the late reply. Thanks for the suggestion of using concrete function! It works for this minimal example and also for my larger network where this issue first appeared :)

Is it in general always recommended to use a concrete function for TFLite conversion to avoid ambiguity? 🤔

[ddavis-2015]

@jianyuzzz

Glad to see you reached a resolution on this issue.

As for using from_concrete_functions, I would do so whenever possible. When converting from a Keras model, "interesting" yet correct operator graphs are generated for TfLite. However, not all such operator graphs work well in TFLM, due to the inherent restrictions on TFLM (ex. no dynamic tensors). The issue #2319 demonstrates this, where a CONV2D operation with dilation > 1 caused the Keras conversion to synthesize additional operations to handle dilation, rather than letting CONV2D do the dilation computation.

[jianyuzzz]

@ddavis-2015

Thanks for the explanation, that clears things up. It's easy to assume TFLM will match TFLite results with the same tflite model, but clearly there are some gotchas. The analyzer and the code snippet in the issue you mentioned are quite efficient to track down where things diverge. Will also use them for the future. Appreciate the help!