`torch.mean` != `tfl.mean`

[maxstrobel]

Description of the bug:

I've found that torch.mean is unwrapped into a tfl.sum & tfl.mul; is there a possibility to avoid this unwrapping & use tfl.mean instead?
This would help me to deploy models more efficiently on an ARM Ethos-U55, which supports the MEAN op on the NPU, but not the SUM.

import ai_edge_torch
import tensorflow as tf
import torch
from torch import nn


class MeanLayer(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return torch.mean(x, dim=1)


mean = MeanLayer()
x_input = torch.randn(1, 16)
with torch.no_grad():
    x_output = mean(x_input)
print(f"Input shape: {x_input.shape}, Output shape: {x_output.shape}")


###
# TFLite quantization + conversion
def _representative_dataset():
    def data_generator():
        for _ in range(3):
            yield [x_input]

    return data_generator


tfl_converter_flags = {
    "optimizations": [tf.lite.Optimize.DEFAULT],
    "target_ops": [tf.lite.OpsSet.TFLITE_BUILTINS_INT8],
    "representative_dataset": _representative_dataset(),
    "inference_input_type": tf.int8,
    "inference_output_type": tf.int8,
    "experimental_enable_resource_variables": True,
}


edge_model = ai_edge_torch.convert(
    mean.eval(),
    (x_input,),
    quant_config=None,
    _ai_edge_converter_flags=tfl_converter_flags,
)
edge_model.export("torch_mean.tflite")

The same thing works directly from Tensorflow:

import tensorflow as tf

inputs = tf.keras.Input(shape=(16,), batch_size=1, dtype=tf.float32)
outputs = tf.keras.layers.Lambda(lambda x: tf.reduce_mean(x, axis=(1,), keepdims=True))(inputs)
model = tf.keras.Model(inputs=inputs, outputs=outputs, name="mean")

model.summary()


converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]
converter.inference_input_type = tf.int8
converter.inference_output_type = tf.int8
converter.representative_dataset = _representative_dataset()
tflite_model = converter.convert()

# Save the model.
with open("tf_mean.tflite", "wb") as f:
    f.write(tflite_model)

Actual vs expected behavior:

torch.mean get converted into tfl.mean.

Any other information you'd like to share?

Python 3.11.10
ai-edge-litert                     1.2.0
ai-edge-quantizer                  0.1.0
ai-edge-torch                      0.4.0
torch                              2.6.0
tensorflow                         2.19.0

[maxstrobel]

@pkgoogle - any opinion on this?

[pkgoogle]

Yeah that's unoptimal, I'm not quite sure why it's happening -- sometimes complex program transformations get lost in the sauce and the fusion of component ops doesn't happen after the transformation or it's hard to reidentify the optimization opportunity.

[maxstrobel]

I took a closer look at the issue.

I think the problem could be that StableHLO has no mean operation (StableHLO ops) defined.

During the export from PyTorch to Tensorflow the operations get replaced - this happens inside exported_programs_to_tflite:

ai-edge-torch/ai_edge_torch/lowertools/_shim.py

Line 47 in 3ca5e5e

def exported_programs_to_tflite(

module {
  func.func public @main(%arg0: tensor<1x16xf32>) -> tensor<1xf32> {
    %0 = call @_aten_mean_3eee5135(%arg0) : (tensor<1x16xf32>) -> tensor<1xf32>
    return %0 : tensor<1xf32>
  }
  func.func private @_aten_mean_3eee5135(%arg0: tensor<1x16xf32>) -> (tensor<1xf32> {jax.result_info = "result"}) {
    %cst = stablehlo.constant dense<0.000000e+00> : tensor<f32>
    %0 = stablehlo.reduce(%arg0 init: %cst) applies stablehlo.add across dimensions = [1] : (tensor<1x16xf32>, tensor<f32>) -> tensor<1xf32>
    %cst_0 = stablehlo.constant dense<1.600000e+01> : tensor<f32>
    %1 = stablehlo.broadcast_in_dim %cst_0, dims = [] : (tensor<f32>) -> tensor<1xf32>
    %2 = stablehlo.divide %0, %1 : tensor<1xf32>
    return %2 : tensor<1xf32>
  }
}

Afterwards, the operations stay separated and aren't merged together.

---

I hope my investigations help here a little bit - do you see a chance to fix this behavior?

[pkgoogle]

Hey @maxstrobel, awesome, that definitely helps! now that I think about it -- it should be converted to a reduction: https://github.com/openxla/stablehlo/blob/main/docs/spec.md#reduce and a reduction w/ these properties should probably be converted to tfl.mean (which we are probably not recognizing properly) -- that's my current hypothesis anyways.

[maxstrobel]

@pkgoogle, if I understand correctly the conversion flow is as follows:

PyTorch
   |
   V
StableHLO    <- Here happens the splitting of MEAN into SUM & MUL
   |
   V
Tensorflow (SavedModel format)
   |
   |                 <- Would we optimize & fuse the operators here?
   V
TFLite

I guess when we have the model in TF, we could somehow merge the operators - or do you see better alternatives?

Do you have an idea how this could work? If you could give me a few pointers where & how you would do this, I would give it a try.

[pkgoogle]

Almost! I recommend using gemini-cli or any other coding agent of your choice to explore https://github.com/llvm/mlir-www and https://github.com/openxla/stablehlo . Basically in this case, you are using a torch model, its exported more or less to StableHLO, but we are using the same mlir/converter infrastructure from TF but as far as I know it's never represented as a TF Saved Model so basically I would just cut that part out. And this is a simplification -- you can think of the PyTorch --> StableHLO part as the first part of a compilation process towards running on hardware, but the StableHLO --> TFLite part as decompiling into a "source graph". As with all compilation and decompilation processes there are many steps and passes -- it gets quite complex. So basically during decompiling there is probably an optimization pass which should probably look for sub-graph patterns which are viable for "fusion" as generally fused ops are more efficient overall than separated ops -- this pass is probably not identifying this opportunity -- but it's pretty hard to say until we dig pretty deep down. So that's why I say probably for now 😄 . Also you can think of PT, TF, and TFL/LRT models as dialects of MLIR, StableHLO is technically a dialect of MLIR as well.