Regression in LSTM TFLite Conversion with Keras 3 (TF 2.17): Generated WHILE-Loop Graph Hangs During TFLite Micro Execution on Eval Board

We are trying to run Keras LSTM using TFlite micro on eval board

Environment
Working:
TensorFlow: 2.15.1
tf.keras: 2.15.0
keras package: 2.15.0
Target: TFLite Micro (eval board)
Converted op: UnidirectionalSequenceLSTM

Broken:
TensorFlow: 2.17.0
tf.keras: 3.15.0 (Keras 3)
keras package: 3.15.0
Target: TFLite Micro (eval board)
Converted op: WHILE

Model configuration
unroll=False (default)
sequence_length=100
4 stacked LSTM layers, hidden_size=128, input_size=32, batch_size=1, dropout=0.0

Problem
We build a 4-layer stacked tf.keras.layers.LSTM model and convert it to TFLite with TFLiteConverter.from_keras_model().

• With TensorFlow 2.15.1 / Keras 2.15.0, the converter fuses the model into UnidirectionalSequenceLSTM ops. Inference on TFLite Micro completes correctly at sequence_length=100.

• With TensorFlow 2.17.0 / Keras 3.15.0 (same model script, only the environment differs), the converter instead emits a graph built around WHILE ops. On TFLite Micro:
  At sequence_length=10, inference completes.
  At sequence_length=100, inference never completes — the debugger stack trace shows the interpreter stuck inside WhileEval → FullyConnected, i.e. the LSTM loop body hangs and never returns.

Standalone code to reproduce the issue -- with keras version 3

import numpy as np
import pathlib
try:
    import tensorflow as tf
except ImportError:
    print("Tensorflow is missing!")
    raise

# Model hyperparameters
INPUT_SIZE = 32
HIDDEN_SIZE = 128
NUM_LAYERS = 4
SEQUENCE_LENGTH = 100
DROPOUT = 0.0
BATCH_SIZE = 1

def build_lstm(
    input_size: int = 32,
    hidden_size: int = 128,
    num_layers: int = 4,
    dropout: float = 0.0,
    sequence_length: int = 100,
    batch_size: int = 1
):
    """
    Build LSTM model with unroll=False (default) to avoid unrolling of LSTM for TFLite conversion.

    Args:
        input_size: Input feature dimension
        hidden_size: LSTM hidden units
        num_layers: Number of stacked LSTM layers
        dropout: Dropout rate between layers
        sequence_length: Fixed sequence length
        batch_size: Fixed batch size

    Returns:
        Keras Model ready for TFLite conversion
    """
    inputs = tf.keras.Input(
        shape=(sequence_length, input_size),
        batch_size=batch_size,
        name='input'
    )

    x = inputs
    for i in range(num_layers):
        x = tf.keras.layers.LSTM(
            hidden_size,
            return_sequences=(i < num_layers - 1),
            dropout=dropout if i < num_layers - 1 else 0.0,
            name=f"lstm_{i+1}"
        )(x)

    outputs = tf.keras.layers.Activation('relu', name='activation')(x)

    return tf.keras.Model(inputs=inputs, outputs=outputs)

def convert_to_tflite(model):
    converter = tf.lite.TFLiteConverter.from_keras_model(model)
    return converter.convert()

def main():
    """Build and convert LSTM model to TFLite"""

    model = build_lstm(
        input_size=INPUT_SIZE,
        hidden_size=HIDDEN_SIZE,
        num_layers=NUM_LAYERS,
        dropout=DROPOUT,
        sequence_length=SEQUENCE_LENGTH,
        batch_size=BATCH_SIZE
    )

    print(f"Input size: {INPUT_SIZE}")
    print(f"Hidden size: {HIDDEN_SIZE}")
    print(f"Num layers: {NUM_LAYERS}")
    print(f"Sequence length: {SEQUENCE_LENGTH}")
    print(f"Batch size: {BATCH_SIZE}")

    test_input = tf.random.normal((1, SEQUENCE_LENGTH, INPUT_SIZE))
    output = model(test_input, training=False)
    print(f"Output size: {output.shape[1]}")

    keras_out_path = pathlib.Path(__file__).parent / f"LSTM_model_seq{SEQUENCE_LENGTH}.keras"
    model.save(keras_out_path)
    print(f"Saved Keras model: {keras_out_path}")

    tflite_model = convert_to_tflite(model)
    out_path = pathlib.Path(__file__).parent / f"LSTM_model_seq{SEQUENCE_LENGTH}.tflite"
    out_path.write_bytes(tflite_model)
    print(f"Saved: {out_path} ({len(tflite_model)/1024:.1f} KB)")

if __name__ == "__main__":
    main()

Question for maintainers
With the same WHILE-loop graph (TF 2.17.0 / Keras 3.15.0, unroll=False) on TFLite Micro:

sequence_length=10 → inference completes successfully.
sequence_length=100 → inference hangs indefinitely inside WhileEval → FullyConnected.
Since it's the same op pattern (WHILE wrapping the LSTM cell body) at both sequence lengths, only the iteration count differs:

• Is there a known iteration-count-dependent issue in the TFLite Micro WhileEval kernel (e.g., a growing/leaking scratch buffer, resource-tensor TfLiteResource reallocation on each iteration, or arena fragmentation) that only manifests once the loop runs enough iterations?

• Is this a genuine infinite loop (loop condition never becomes false) that happens to "run out of time" before we notice at 10 iterations, or is it actually a soft hang (e.g., an unbounded scratch-memory allocation per iteration causing the arena/allocator to effectively stall as it grows)?

• Is there a maximum supported/tested iteration count for WHILE-based LSTM on TFLite Micro, and is UnidirectionalSequenceLSTM fusion the only supported path for long sequences?

This distinction matters for us in deciding whether to wait for a converter/runtime fix vs. permanently restructuring our model to avoid WHILE entirely.