Possibly Erroneous FLOPS Calculation

### System information
Operating System: Linux Ubuntu 22.04
JAX Version: 0.4.17
jaxlib Version: 0.4.14
Flax Version: 0.7.5
Python Version: 3.10.12
GPU: T4
CUDA: 11.8

### Problem you have encountered:
I applied a Conv block with 2 features with (2,2) kernel size, no bias, "same" padding, (1,1) strides to a (1,4,4,1) array both in flax and keras, then calculated flops used for each. Keras flop calculation would return 256 which makes sense (2 kernels * 16 output cells * (4 mul + 4 sum) = 256). However, both the JAX internal flop calculation tool and the flax flop calculation from the tabulate function return 196 as flops used. How can that be? If we were to dissect the number 196 by dividing it by 2 kernels and 16 output cells, we would get 6.125 flops per cell. The only way I think this would be possible is 4 mul and 2 1/8th add per cell with some sort of vectorization optimization going on? This discrepancy between keras and flax flop count doesn't occur if the padding is "valid".

### Steps to reproduce:
You can find my short code in the link:
https://colab.research.google.com/drive/1ClNx5FWMIvf_H_jGVV4Z2T1_9J0ZmNlF?usp=sharing