Wrong quantized_dimension (axis) when "per-channel" quantization
Hi @CarlosNacher ,

Can you give me the full code to replicate this issue. The library imports and some other things are not clear . I will be able to debug quicker if you provide me the above details.  Hey @sawantkumar ,

For sure! I have uploaded the code to a [Colab](https://colab.research.google.com/drive/17StLOhZJplMbZ2tew0UF85qXSUWd22vY?usp=sharing) file so that you can fully reproduce it. I have also written some comments in the notebook explaining the actual behaviour and the expected one.

The link to the data you will need to fully reproduce (the SavedModel and calibration data) is here: https://we.tl/t-MUNhBovDZr

Thank you so much for your response and for taking the time to help me!
 Hi @CarlosNacher ,

Sorry for the late response, but the link for your savedModel seems to have expired. can you please provide another link to the model and data.  Hi @sawantkumar ,

For sure, here: https://we.tl/t-iUsgchLLL7

Thank you. Hi @sawantkumar , could you please give me some feedback aboout this? The links only last 7 days. If you need it, I can send it in a different format.

Thanks! Hi @CarlosNacher ,

I replicated your script and i got the tflite and i inspected it and confirmed that its quantizing the 0th dimension.  

 ````
dtype=float32),` 'zero_points': array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
       0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]), 'quantized_dimension': 0}
Tensor Name: tfl.pseudo_qconst5``
``` Hi @CarlosNacher ,

I will raise a PR to fix your issue and I will update you. Thank you Hi @sawantkumar ,

Okay, thank you so much! It's great that we can all contribute together. I look forward to your response. Hi @CarlosNacher ,

Sure thing. However can you please confirm if this bug can be reproduced in tf 2.16.0 or tf-nightly .  Hey @sawantkumar ! I confirm it, I have replicated using tf 2.16 and tf-nightly (2.17.0.dev20240528), the output trace of the code execution is in the Colab I attached. Thank you. Hello @sawantkumar , 

Any update on this? If there is a PR, could I havee the link to keep updated? 

Thank you in advance! Hi @CarlosNacher, this issue may be resolved by following a different work flow. 

Have you tried using [AI-Edge-Torch](https://github.com/google-ai-edge/ai-edge-torch)?, you can find more information here: [googleblog](https://developers.googleblog.com/en/ai-edge-torch-high-performance-inference-of-pytorch-models-on-mobile-devices/).

I have actually created a simple script for converting a CV model here:

```py
import torch
import torchvision
import ai_edge_torch

mobilenet_model = torchvision.models.mobilenet_v3_small()
sample_inputs = (torch.randn(1, 3, 224, 224),)

edge_model = ai_edge_torch.convert(mobilenet_model.eval(), sample_inputs)
edge_model.export("mobilenet_v3_small.tflite")
```

If you want to, you can actually try visualizing the result in [model-explorer](https://github.com/google-ai-edge/model-explorer) as well.

Please try them out and let us know if this resolves your issue. If you still need further help, feel free to open a new issue at the respective repo. Hi @pkgoogle 

I didn’t know that library, thank you so much for the information! I’ll check it out and let you know, however my final application of interest is to run the model in coral TPU USB accelerator, do you know if with the tool you comment it is achievable? 
On the Readme, CPU and GPU support is mentioned but not TPU. As far I know, the only thing you need to run your .tflite on TPU is that it is FULLY INT8, so I suppose that it’ll work, but the no mention of TPU confuses me.

Can you give me feedback about that? Anyway, again, it’ll try, and many thanks!
### 1. System information

- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Microsoft Windows 11 Home 10.0.22631 version 22631 compilation 
- TensorFlow installation (pip package or built from source): `pip install tensorflow==2.15.0`
- TensorFlow library (version, if pip package or github SHA, if built from source): 2.15.0

### 2. Code

```
# Create callib data
input_op_name = 'input'
numpy_file_path = 'calibration_data.npy'
calib_data = np.load(numpy_file_path)
data_count = calib_data.shape[0]
mean = [[[[0, 0, 0]]]]
std = [[[[1, 1, 1]]]]
calib_data_dict = {}
calib_data_dict[input_op_name] = \
    [
    calib_data.copy(),
    mean,
    std,
    ]

# representative_dataset_gen
def representative_dataset_gen():
    for idx in range(data_count):
        yield_data_dict = {}
        for model_input_name in [input_op_name]:
            calib_data, mean, std = calib_data_dict[model_input_name]
            normalized_calib_data: np.ndarray = (calib_data[idx] - mean) / std
            yield_data_dict[model_input_name] = tf.cast(tf.convert_to_tensor(normalized_calib_data), tf.float32)
        yield yield_data_dict


###### CONVERT
converter.optimizations = [tf.lite.Optimize.DEFAULT]
converter.target_spec.supported_ops = [
    tf.lite.OpsSet.TFLITE_BUILTINS_INT8,
    tf.lite.OpsSet.SELECT_TF_OPS,
]
converter._experimental_disable_per_channel = False # to perform "per-channel" (it's the default value though)
converter.representative_dataset = representative_dataset_gen
converter.inference_input_type = "int8"
converter.inference_output_type = "int8"

TFLITE_FILEPATH = "model_full_int.tflite"
with open(TFLITE_FILEPATH , 'wb') as w:
    w.write(tflite_model)

##### INFER
# Initialize the interpreter
interpreter = tf.lite.Interpreter(
    model_path=str(TFLITE_FILEPATH"), 
    # experimental_preserve_all_tensors=True
    )
interpreter.allocate_tensors()

input_details = interpreter.get_input_details()[0]
output_details = interpreter.get_output_details()[0]

test_image = pre_processed_input.copy() # One image I have loaded (shape 1, 1024, 608, 3)

# Check if the input type is quantized, then rescale input data to int8
if input_details['dtype'] in [np.int8, np.uint8, np.int16]:
    input_scale, input_zero_point = input_details["quantization"]
    test_image = test_image / input_scale + input_zero_point
test_image = np.expand_dims(test_image, axis=0).astype(input_details["dtype"])


interpreter.set_tensor(input_details["index"], test_image)
interpreter.invoke()
tflite_inference = interpreter.get_tensor(output_details["index"])[0]

if output_details["dtype"] in [np.int8, np.uint8, np.int16]:
    output_scale, output_zero_point = output_details['quantization']
    print("Output scale:", output_scale)
    print("Output zero point:", output_zero_point)
    print()
    tflite_inference = output_scale * (tflite_inference.astype(np.float32) - output_zero_point)

```



### 3. Failure after conversion
The conversion is successful, but the generated model is wrong: **the inferred output (`tflite_inference`) is all the same value in the fully int8 quantized model and I think it could be because quantization per-channel is performing per-batchaxis instead**:

If I inspect the .tflite model, I found that in each operation (for example convolution) the number of `scale` and `zero_point` values are the same as number of axis=0 (batch i.e 256) not the number of channels (axis=-1 i.e 512)

![image](https://github.com/tensorflow/tensorflow/assets/85699494/bf57d6ea-a2b4-41bc-a43b-9809eb393cde)
...
![image](https://github.com/tensorflow/tensorflow/assets/85699494/519b7a9e-0469-40dc-bfbf-755af63d0706)

**I miss one `quantized_dimension` parameter, as stated in docs:** https://www.tensorflow.org/lite/performance/quantization_spec#per-axis_vs_per-tensor 