Support dynamic shape output for CompiledModel

[leeroyka]

LiteRT v2.1.1 CompiledModel C++ Android

The model has a dynamic shape of the input and output tensors.
When I change the size of the input tensor, the output shape does not change.

Example:

my model has:

Input:
  0: name=data, shape=[ 1 -1 -1 3], dtype=float32

Output:
  0: name=out_0, shape=[ 1 -1 -1 18], dtype=float32
  1: name=out_1, shape=[ 1 -1 -1 18], dtype=float32
  2: name=out_2, shape=[ 1 -1 -1 18], dtype=float32
  3: name=out_3, shape=[ 1 -1 -1 18], dtype=float32

c++ code:

    LITERT_ABORT_IF_ERROR(compiled_model.ResizeInputTensor("data", {1, 1088, 1920, 3}));
    LITERT_ASSIGN_OR_ABORT(auto input_buffers,
                         compiled_model.CreateInputBuffers());
    LITERT_ASSIGN_OR_ABORT(auto output_buffers,
                         compiled_model.CreateOutputBuffers());

    auto input_buffer_size = input_buffers[0].Size();
    LOGI("INPUT_BUFFER_SIZE: %zu", *input_buffer_size);
    LOGI("OUTPUT_BUFFER_SIZE: %zu", output_buffers.size());
    for (size_t i = 0; i < output_buffers.size(); ++i) {
        auto output_buffers_size = output_buffers[i].Size();
        LOGI("Output buffer %zu size: %zu", i, *output_buffers_size);
    }

logs:

INPUT_BUFFER_SIZE: 25067520
OUTPUT_BUFFER_SIZE: 4
Output buffer 0 size: 72
Output buffer 1 size: 72
Output buffer 2 size: 72
Output buffer 3 size: 72

but for my models output shapes should be:

  out_0 shape=(1, 68, 120, 18)
  out_1 shape=(1, 136, 240, 18)
  out_2 shape=(1, 17, 30, 18)
  out_3 shape=(1, 34, 60, 18) 

[aravindmurali22]

Hi @leeroyka,

Thank you for reporting this.

I've analyzed the CompiledModel C++ implementation, and your usage looks correct. Specifically, CreateOutputBuffers() internally calls GetOutputTensorLayouts(..., update_allocation=true), which should trigger shape propagation and resize the output tensors in the underlying runtime before querying for their buffer requirements.

However, the fact that you are getting 72 bytes (which matches 1 * 1 * 1 * 18 * sizeof(float)) instead of the expected size suggests that GetOutputBufferRequirements is returning stale information corresponding to the initial static shape (or the default -1 placeholders treated as 1), ignoring the propagated shapes from ResizeInputTensor.

To help us pinpoint where the update is failing, could you please check the following:

1. Are you running this on CPU, GPU, or NPU? (Pass HwAccelerators::kCpu / kGpu / kNpu to CompiledModel::Create?)
2. Could you print the dimensions returned by GetOutputTensorType(output_name).Layout().Dimensions() after the resize but before creating buffers? This will tell us if the logical shape update succeeded even if the buffer requirement failed.

   // After ResizeInputTensor...
   auto tensor_type = compiled_model.GetOutputTensorType("out_0");
   auto dims = tensor_type->Layout().Dimensions();
   LOGI("Output 0 dims: %d, %d, %d, %d", dims[0], dims[1], dims[2], dims[3]);
   

3. Is it possible to share a minimal reproduction model (tflite file) that exhibits this behavior?

If GetOutputTensorType returns the correct new dimensions but CreateOutputBuffers creates a tiny buffer, it confirms a bug in GetCompiledModelOutputBufferRequirements not respecting the updated state.

[snnn]

You need to check the model if shape signature was correctly set. In the model file each tensor has a 'shape' and a 'shape_signature' field. The shape array should not have negative numbers, while the shape_signature array can have '-1' to represent unknown(other negative numbers are not allowed).

At runtime when a model is loaded, in tflite/core/interpreter_builder.cc the runtime code just copies FlatBuffer shape_signature into dims_signature. Then, when running inference, for each input buffer in Run(...), RegisterBuffer(...) is called. The RegisterBuffer(...) extracts the buffer layout shape and calls InputTensorNeedsResize. The InputTensorNeedsResize function decides:

• false if shape already matches
• error if no dims_signature or no dynamic dim (-1)
• error if rank/static dims mismatch
• true if compatible dynamic resize is needed

Then, if resize is needed, it calls TFLite SignatureRunner::ResizeInputTensor.

So, please be sure input shape does not have negative numbers and shape_signature is correctly set.