Have you checked if is it still related to:

  

      github.com/tensorflow/tensorflow
  

  
    
  
	  
  

  
    
      Softmax activations don't get converted to Softmax TFLite operator if ndim > 2 7
    

    
      
        opened 
        
          
        
        Dec 18, 2019
      
      

        
          closed 
        
          
        
        Mar 6, 2021
      
        

      
        
          
          ghop02
        
      
    

    
        type:bug
        comp:lite
        TF 2.1
    
  


  
    **System information**
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04…): MacOS 10.15.1
- TensorFlow installed from (source or binary): binary
- TensorFlow version (or github SHA if from source): tf-nightly==2.1.0.dev20191203

**Command used to run the converter or code if you’re using the Python API**

```
import pathlib

inpt = tf.keras.layers.Input(shape=[256, 256, 3])
out = tf.keras.layers.Lambda(lambda x: tf.keras.activations.softmax(x))(inpt)
out = tf.keras.layers.Lambda(lambda x: tf.nn.softmax(x))(out)
model = tf.keras.Model(inpt, out)

converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()
pathlib.Path('out.tflite').write_bytes(tflite_model)
```

**Failure details**
![image](https://user-images.githubusercontent.com/1422280/71025023-3904e880-20d4-11ea-95fb-29cfea49a44d.png)
This graph shows the difference between the different softmax methods.  When using `tf.keras.activations.softmax`, there is [code](https://github.com/tensorflow/tensorflow/blob/v2.1.0-rc1/tensorflow/python/keras/activations.py#L43-L79) with a workaround for multiple dimensions.  It looks like this was written before the tensorflow op had multi-dimension support.