[RNN] Keras LSTM converted to "While" OPs with hidden states manipulation - TFLite
Hi @Doomski99, the UnidirectionalSequenceLSTM op may not allow multiple outputs, the body subgraph of the While op might effectively calculate the same thing as the UnidirectionalSequenceLSTM op. You said this op fails to run on your target platform ... what's your target platform? and what's the error? is there a log or any error message? Hello and thank you for your reply @pkgoogle .

> the body subgraph of the While op might effectively calculate the same thing as the UnidirectionalSequenceLSTM op

I understand but can you confirm that the performance is the same between the two operators? Could it be that the UnidirectionalSequenceLSTM op has a better optimized implementation?

> what's your target platform? and what's the error? is there a log or any error message?

Cadence's HiFi 5 DSP. I'm using a simulation environment. The model with While stops after launching with a "HALTED" message unlike the second model which runs perfectly.

> UnidirectionalSequenceLSTM op may not allow multiple outputs

Can I get a confirmation?  Hi @Doomski99,

I was able to create, load, and run both models on colab: https://colab.sandbox.google.com/gist/pkgoogle/7cd581019a15c5db26d9071e3f968294/62775.ipynb .... do you have any limitations or controls in using your simulated environment? Is there any way to relax those constraints? Just judging from the message and my results, I think that message was produced by your environment rather from TFLite (I would've expected an exception if it was thrown from TFLite)

> I understand but can you confirm that the performance is the same between the two operators? Could it be that the UnidirectionalSequenceLSTM op has a better optimized implementation?

You can try (maybe in a general linux system first) to benchmark them: https://github.com/tensorflow/tensorflow/tree/master/tensorflow/lite/tools/benchmark

> Can I get a confirmation?

I can't confirm, here's the code: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/lite/kernels/unidirectional_sequence_lstm.cc It probably depends on what that while body subgraph actually is, i.e. the Op might actually be inside there.

 Thank you for your help @pkgoogle !

Before digging through my DSP simulation environment, I've successfully set up the benchmark tool you linked in WSL2 and ran it through my models with the following command:

`bazel-bin/tensorflow/lite/tools/benchmark/benchmark_model   --graph=1x_LSTM_hidden_64_float32.tflite   --num_threads=1 --enable_op_profiling=true`.

The one with "UnidirectionalSequenceLstm" is consistently 35% faster. Here's the output of both benchmarks:

First (UnidirectionalSequenceLstm) model:

```
bazel-bin/tensorflow/lite/tools/benchmark/benchmark_model   --graph=1x_LSTM_64_float32.tflite   --num_threads=1 --enable_op_profiling=true
INFO: STARTING!
INFO: Log parameter values verbosely: [0]
INFO: Num threads: [1]
INFO: Graph: [1x_LSTM_64_float32.tflite]
INFO: Enable op profiling: [1]
INFO: #threads used for CPU inference: [1]
INFO: Loaded model 1x_LSTM_64_float32.tflite
INFO: Created TensorFlow Lite XNNPACK delegate for CPU.
INFO: The input model file size (MB): 0.330964
INFO: Initialized session in 0.63ms.
INFO: Running benchmark for at least 1 iterations and at least 0.5 seconds but terminate if exceeding 150 seconds.
^[[AINFO: count=132 first=3678 curr=3790 min=3270 max=5473 avg=3751.66 std=334

INFO: Running benchmark for at least 50 iterations and at least 1 seconds but terminate if exceeding 150 seconds.
INFO: count=261 first=3937 curr=3959 min=3190 max=6333 avg=3790.92 std=358

INFO: Inference timings in us: Init: 630, First inference: 3678, Warmup (avg): 3751.66, Inference (avg): 3790.92
INFO: Note: as the benchmark tool itself affects memory footprint, the following is only APPROXIMATE to the actual memory footprint of the model at runtime. Take the information at your discretion.
INFO: Memory footprint delta from the start of the tool (MB): init=5.16016 overall=7.46484
INFO: Profiling Info for Benchmark Initialization:
============================== Run Order ==============================
                                     [node type]          [first]        [avg ms]            [%]          [cdf%]          [mem KB]      [times called]  [Name]
                         ModifyGraphWithDelegate            0.104           0.104        83.871%         83.871%             0.000              1       ModifyGraphWithDelegate/0
                                 AllocateTensors            0.020           0.020        16.129%        100.000%             0.000              1       AllocateTensors/0

============================== Top by Computation Time ==============================
                                     [node type]          [first]        [avg ms]            [%]          [cdf%]          [mem KB]      [times called]  [Name]
                         ModifyGraphWithDelegate            0.104           0.104        83.871%         83.871%             0.000              1       ModifyGraphWithDelegate/0
                                 AllocateTensors            0.020           0.020        16.129%        100.000%             0.000              1       AllocateTensors/0

Number of nodes executed: 2
============================== Summary by node type ==============================
                                     [Node type]          [count]         [avg ms]          [avg %]         [cdf %]       [mem KB]      [times called]
                         ModifyGraphWithDelegate                1            0.104          83.871%         83.871%          0.000              1
                                 AllocateTensors                1            0.020          16.129%        100.000%          0.000              1

Timings (microseconds): count=1 curr=124
Memory (bytes): count=0
2 nodes observed



INFO: Operator-wise Profiling Info for Regular Benchmark Runs:
============================== Run Order ==============================
                                     [node type]          [first]        [avg ms]            [%]          [cdf%]          [mem KB]      [times called]  [Name]
                                  RANDOM_UNIFORM            0.003           0.003         0.083%          0.083%             0.000              1       [simple_model/random_uniform/RandomUniform]:0
                                  RANDOM_UNIFORM            0.003           0.003         0.070%          0.153%             0.000              1       [simple_model/random_uniform_1/RandomUniform]:1
                    UNIDIRECTIONAL_SEQUENCE_LSTM            3.912           3.769        99.648%         99.801%             0.000              1       [tfl.unidirectional_sequence_lstm]:2
                  Fully Connected (NC, F32) GEMM            0.007           0.008         0.199%        100.000%             0.000              1       Delegate/Fully Connected (NC, F32) GEMM:0

============================== Top by Computation Time ==============================
                                     [node type]          [first]        [avg ms]            [%]          [cdf%]          [mem KB]      [times called]  [Name]
                    UNIDIRECTIONAL_SEQUENCE_LSTM            3.912           3.769        99.648%         99.648%             0.000              1       [tfl.unidirectional_sequence_lstm]:2
                  Fully Connected (NC, F32) GEMM            0.007           0.008         0.199%         99.847%             0.000              1       Delegate/Fully Connected (NC, F32) GEMM:0
                                  RANDOM_UNIFORM            0.003           0.003         0.083%         99.930%             0.000              1       [simple_model/random_uniform/RandomUniform]:0
                                  RANDOM_UNIFORM            0.003           0.003         0.070%        100.000%             0.000              1       [simple_model/random_uniform_1/RandomUniform]:1

Number of nodes executed: 4
============================== Summary by node type ==============================
                                     [Node type]          [count]         [avg ms]          [avg %]         [cdf %]       [mem KB]      [times called]
                    UNIDIRECTIONAL_SEQUENCE_LSTM                1            3.768          99.683%         99.683%          0.000              1
                  Fully Connected (NC, F32) GEMM                1            0.007           0.185%         99.868%          0.000              1
                                  RANDOM_UNIFORM                2            0.005           0.132%        100.000%          0.000              2

Timings (microseconds): count=261 first=3925 curr=3943 min=3186 max=6312 avg=3781.84 std=355
Memory (bytes): count=0
4 nodes observed
```
Second (While) model:

```
bazel-bin/tensorflow/lite/tools/benchmark/benchmark_model   --graph=1x_LSTM_hidden_64_float32.tflite   --num_threads=1 --enable_op_profiling=true
INFO: STARTING!
INFO: Log parameter values verbosely: [0]
INFO: Num threads: [1]
INFO: Graph: [1x_LSTM_hidden_64_float32.tflite]
INFO: Enable op profiling: [1]
INFO: #threads used for CPU inference: [1]
INFO: Loaded model 1x_LSTM_hidden_64_float32.tflite
INFO: Created TensorFlow Lite XNNPACK delegate for CPU.
WARNING: Attempting to use a delegate that only supports static-sized tensors with a graph that has dynamic-sized tensors (tensor#10 is a dynamic-sized tensor).
INFO: The input model file size (MB): 0.514668
INFO: Initialized session in 0.609ms.
INFO: Running benchmark for at least 1 iterations and at least 0.5 seconds but terminate if exceeding 150 seconds.
INFO: count=107 first=5292 curr=4749 min=4243 max=5357 avg=4657.9 std=246

INFO: Running benchmark for at least 50 iterations and at least 1 seconds but terminate if exceeding 150 seconds.
INFO: Warning: Dropping ProfileBuffer event.
INFO: count=191 first=4983 curr=5046 min=4295 max=6180 avg=4918.73 std=320

INFO: Inference timings in us: Init: 609, First inference: 5292, Warmup (avg): 4657.9, Inference (avg): 4918.73
INFO: Note: as the benchmark tool itself affects memory footprint, the following is only APPROXIMATE to the actual memory footprint of the model at runtime. Take the information at your discretion.
INFO: Memory footprint delta from the start of the tool (MB): init=4.94141 overall=10.2656
INFO: Profiling Info for Benchmark Initialization:
Primary graph profile:
============================== Run Order ==============================
                                     [node type]          [first]        [avg ms]            [%]          [cdf%]          [mem KB]      [times called]  [Name]
                                 AllocateTensors            0.018           0.018        19.355%         19.355%             0.000              1       AllocateTensors/0
                         ModifyGraphWithDelegate            0.075           0.075        80.645%        100.000%             0.000              1       ModifyGraphWithDelegate/0

============================== Top by Computation Time ==============================
                                     [node type]          [first]        [avg ms]            [%]          [cdf%]          [mem KB]      [times called]  [Name]
                         ModifyGraphWithDelegate            0.075           0.075        80.645%         80.645%             0.000              1       ModifyGraphWithDelegate/0
                                 AllocateTensors            0.018           0.018        19.355%        100.000%             0.000              1       AllocateTensors/0

Number of nodes executed: 2
============================== Summary by node type ==============================
                                     [Node type]          [count]         [avg ms]          [avg %]         [cdf %]       [mem KB]      [times called]
                         ModifyGraphWithDelegate                1            0.075          80.645%         80.645%          0.000              1
                                 AllocateTensors                1            0.018          19.355%        100.000%          0.000              1

Timings (microseconds): count=1 curr=93
Memory (bytes): count=0
2 nodes observed

Subgraph (index: 1) profile:
============================== Run Order ==============================
                                     [node type]          [first]        [avg ms]            [%]          [cdf%]          [mem KB]      [times called]  [Name]
                                 AllocateTensors            0.007           0.007       100.000%        100.000%             0.000              1       AllocateTensors/1

============================== Top by Computation Time ==============================
                                     [node type]          [first]        [avg ms]            [%]          [cdf%]          [mem KB]      [times called]  [Name]
                                 AllocateTensors            0.007           0.007       100.000%        100.000%             0.000              1       AllocateTensors/1

Number of nodes executed: 1
============================== Summary by node type ==============================
                                     [Node type]          [count]         [avg ms]          [avg %]         [cdf %]       [mem KB]      [times called]
                                 AllocateTensors                1            0.007         100.000%        100.000%          0.000              1

Subgraph (index: 1) Timings (microseconds): count=1 curr=7
Memory (bytes): count=0
1 nodes observed

Subgraph (index: 2) profile:
============================== Run Order ==============================
                                     [node type]          [first]        [avg ms]            [%]          [cdf%]          [mem KB]      [times called]  [Name]
                                 AllocateTensors            0.025           0.016       100.000%        100.000%             0.000              2       AllocateTensors/2

============================== Top by Computation Time ==============================
                                     [node type]          [first]        [avg ms]            [%]          [cdf%]          [mem KB]      [times called]  [Name]
                                 AllocateTensors            0.025           0.016       100.000%        100.000%             0.000              2       AllocateTensors/2

Number of nodes executed: 1
============================== Summary by node type ==============================
                                     [Node type]          [count]         [avg ms]          [avg %]         [cdf %]       [mem KB]      [times called]
                                 AllocateTensors                1            0.032         100.000%        100.000%          0.000              2

Subgraph (index: 2) Timings (microseconds): count=1 curr=32
Memory (bytes): count=0
1 nodes observed



INFO: Operator-wise Profiling Info for Regular Benchmark Runs:
Primary graph profile:
============================== Run Order ==============================
                                     [node type]          [first]        [avg ms]            [%]          [cdf%]          [mem KB]      [times called]  [Name]
                                  RANDOM_UNIFORM            0.003           0.003         0.067%          0.067%             0.000              1       [simple_model/random_uniform/RandomUniform]:0
                                       TRANSPOSE            0.027           0.031         0.628%          0.695%             0.000              1       [transpose]:1
                                           WHILE            4.931           4.863        99.305%        100.000%             0.000              1       [while, while1, while2, while3, while4, Unknown]:2

============================== Top by Computation Time ==============================
                                     [node type]          [first]        [avg ms]            [%]          [cdf%]          [mem KB]      [times called]  [Name]
                                           WHILE            4.931           4.863        99.305%         99.305%             0.000              1       [while, while1, while2, while3, while4, Unknown]:2
                                       TRANSPOSE            0.027           0.031         0.628%         99.933%             0.000              1       [transpose]:1
                                  RANDOM_UNIFORM            0.003           0.003         0.067%        100.000%             0.000              1       [simple_model/random_uniform/RandomUniform]:0

Number of nodes executed: 3
============================== Summary by node type ==============================
                                     [Node type]          [count]         [avg ms]          [avg %]         [cdf %]       [mem KB]      [times called]
                                           WHILE                1            4.862          99.326%         99.326%          0.000              1
                                       TRANSPOSE                1            0.030           0.613%         99.939%          0.000              1
                                  RANDOM_UNIFORM                1            0.003           0.061%        100.000%          0.000              1

Timings (microseconds): count=191 first=4961 curr=5027 min=4281 max=6154 avg=4896.75 std=315
Memory (bytes): count=0
3 nodes observed

Subgraph (index: 1) profile:
============================== Run Order ==============================
                                     [node type]          [first]        [avg ms]            [%]          [cdf%]          [mem KB]      [times called]  [Name]
                                            LESS            0.000           0.000        61.756%         61.756%             0.000             37       [simple_model/dense/Tensordot/MatMul1]:0
                                 AllocateTensors            0.003           0.004        38.244%        100.000%             0.000              1       AllocateTensors/1

============================== Top by Computation Time ==============================
                                     [node type]          [first]        [avg ms]            [%]          [cdf%]          [mem KB]      [times called]  [Name]
                                 AllocateTensors            0.003           0.004        38.244%         38.244%             0.000              1       AllocateTensors/1
                                            LESS            0.000           0.000        61.756%        100.000%             0.000             37       [simple_model/dense/Tensordot/MatMul1]:0

Number of nodes executed: 2
============================== Summary by node type ==============================
                                     [Node type]          [count]         [avg ms]          [avg %]         [cdf %]       [mem KB]      [times called]
                                            LESS                1            0.005          62.500%         62.500%          0.000             37
                                 AllocateTensors                1            0.003          37.500%        100.000%          0.000              1

Subgraph (index: 1) Timings (microseconds): count=191 first=9 curr=7 min=3 max=61 avg=9.24084 std=5
Memory (bytes): count=0
2 nodes observed

Subgraph (index: 2) profile:
============================== Run Order ==============================
                                     [node type]          [first]        [avg ms]            [%]          [cdf%]          [mem KB]      [times called]  [Name]
                                 AllocateTensors            0.024           0.010         0.263%          0.263%             0.000              1       AllocateTensors/2
                                             ADD            0.001           0.000         0.137%          0.400%             0.000             37       [transpose_1]:0                                 FULLY_CONNECTED            0.018           0.017        16.809%         17.208%             0.000             37       [StatefulPartitionedCall:0]:1
                                          GATHER            0.001           0.001         0.785%         17.993%             0.000             37       [StatefulPartitionedCall:1]:2
                                 FULLY_CONNECTED            0.091           0.068        65.937%         83.930%             0.000             37       [Unknown]:3
                                             ADD            0.001           0.001         0.650%         84.580%             0.000             37       [Unknown]:4
                                             ADD            0.001           0.001         0.623%         85.203%             0.000             37       [Unknown]:5
                                           SPLIT            0.001           0.001         0.639%         85.841%             0.000             37       [Unknown, Unknown, Unknown, Unknown]:6
                                        LOGISTIC            0.001           0.001         0.545%         86.386%             0.000             36       [Unknown]:7
                                        LOGISTIC            0.001           0.000         0.397%         86.783%             0.000             36       [Unknown]:8
                                             MUL            0.001           0.000         0.416%         87.199%             0.000             36       [Unknown]:9
                                        LOGISTIC            0.001           0.000         0.419%         87.618%             0.000             36       [Unknown]:10
                                             ADD            0.000           0.000         0.147%         87.765%             0.000             36       [Unknown]:24
                                   CONCATENATION            0.003           0.004         3.537%         91.302%             0.000             36       [Unknown]:23
                                         RESHAPE            0.000           0.000         0.046%         91.348%             0.000             36       [Unknown]:22
                                           SLICE            0.008           0.004         3.847%         95.195%             0.000             36       [Unknown]:21
                                            TANH            0.000           0.000         0.401%         95.596%             0.000             36       [Unknown]:11
                                             MUL            0.001           0.000         0.265%         95.861%             0.000             36       [Unknown]:12
                                             ADD            0.000           0.000         0.221%         96.082%             0.000             36       [Unknown]:13
                                            TANH            0.000           0.000         0.360%         96.442%             0.000             36       [Unknown]:14
                                             MUL            0.001           0.000         0.261%         96.703%             0.000             36       [Unknown]:15
                                         RESHAPE            0.000           0.000         0.059%         96.762%             0.000             36       [Unknown]:16
                                   CONCATENATION            0.001           0.000         0.282%         97.044%             0.000             36       [Unknown]:17
                                           SLICE            0.001           0.003         2.774%         99.818%             0.000             36       [Unknown]:18
                                         RESHAPE            0.000           0.000         0.049%         99.867%             0.000             36       [Unknown]:19
                                   CONCATENATION            0.000           0.000         0.133%        100.000%             0.000             36       [Unknown]:20

============================== Top by Computation Time ==============================
                                     [node type]          [first]        [avg ms]            [%]          [cdf%]          [mem KB]      [times called]  [Name]
                                 FULLY_CONNECTED            0.091           0.068        65.937%         65.937%             0.000             37       [Unknown]:3
                                 FULLY_CONNECTED            0.018           0.017        16.809%         82.745%             0.000             37       [StatefulPartitionedCall:0]:1
                                 AllocateTensors            0.024           0.010         0.263%         83.008%             0.000              1       AllocateTensors/2
                                           SLICE            0.008           0.004         3.847%         86.855%             0.000             36       [Unknown]:21
                                   CONCATENATION            0.003           0.004         3.537%         90.393%             0.000             36       [Unknown]:23
                                           SLICE            0.001           0.003         2.774%         93.167%             0.000             36       [Unknown]:18
                                          GATHER            0.001           0.001         0.785%         93.951%             0.000             37       [StatefulPartitionedCall:1]:2
                                             ADD            0.001           0.001         0.650%         94.601%             0.000             37       [Unknown]:4
                                           SPLIT            0.001           0.001         0.639%         95.240%             0.000             37       [Unknown, Unknown, Unknown, Unknown]:6
                                             ADD            0.001           0.001         0.623%         95.863%             0.000             37       [Unknown]:5

Number of nodes executed: 26
============================== Summary by node type ==============================
                                     [Node type]          [count]         [avg ms]          [avg %]         [cdf %]       [mem KB]      [times called]
                                 FULLY_CONNECTED                2            3.157          83.035%         83.035%          0.000             74
                                           SLICE                2            0.251           6.602%         89.637%          0.000             72
                                   CONCATENATION                3            0.150           3.945%         93.582%          0.000            108
                                             ADD                5            0.065           1.710%         95.292%          0.000            183
                                        LOGISTIC                3            0.050           1.315%         96.607%          0.000            108
                                             MUL                3            0.034           0.894%         97.501%          0.000            108
                                          GATHER                1            0.029           0.763%         98.264%          0.000             37
                                            TANH                2            0.028           0.736%         99.001%          0.000             72
                                           SPLIT                1            0.024           0.631%         99.632%          0.000             37
                                 AllocateTensors                1            0.010           0.263%         99.895%          0.000              1
                                         RESHAPE                3            0.004           0.105%        100.000%          0.000            108

Subgraph (index: 2) Timings (microseconds): count=191 first=3863 curr=3956 min=3340 max=4844 avg=3816.78 std=256
Memory (bytes): count=0
26 nodes observed
```

It's interesting that we can see the operations happening inside the While loop through this tool. 

> It probably depends on what that while body subgraph actually is, i.e. the Op might actually be inside there.

I guess the Op isn't there after all. It seems that it's splitting the LSTM operation.  What I find interesting is that, in the second model, the "subgraph index = 2" that is mentioned inside the While loop is averaging at 3.8ms, which is close to that of the UnidirectionalSequenceLSTM, but  the actual while loop is averaging at 4.8ms, more than the "subgraph index = 2". Not sure if that's relevant or not.

Anyways, I've also tried another LSTM model of mine that is dynamically quantized, and I'm getting a 40% speedup. I don't think we can justify the speedup for simply not retaining the hidden_states, right?

Performance is critical for my task, how do you think we should proceed? I'll be waiting for your input.
 So the converter is in some sense an arbitrary program optimizer, since TF is Turing complete, it probably doesn't handle all cases to complete optimum, there probably exists a conversion which is more optimal (or maybe we need to change the source code/that op to more output those hidden states natively rather than the deconstruction/reconstruction you saw). This will likely be a hard problem to solve generally and you say it's critical for your application. We will get to it eventually but we have to consider all issues and how many issues/users each potential change can help with and their relative impact as well -- so it might take a while to get to this one. That being said, since it is impactful for you and now that you are touching the innards a bit you can try your best to see if you can introduce a PR for this yourself.

Let's change this into a feature/performance issue and I can direct you a little bit on what to look at but you'll likely have to do a lot of digging for yourself. Learn some MLIR https://mlir.llvm.org/ (A very general compiler framework ... a compiler translates higher level code to lower level code, but some of the same abstractions can be used to .. side translate/optimize such as TF -> TFL, optimization can be viewed as lowering to a representation more efficient for your specific configuration/hardware)

Build from source with debug info, especially the tflite_convert CLT: https://www.tensorflow.org/lite/models/convert/convert_models#command_line_tool, use lldb/gdb with it and dig into how your program gets transformed and see how we can do the optimization better

This README is probably useful: https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/lite/README.md

Some of the binaries here might be useful:
https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/mlir/lite/BUILD

Hi @majiddadashi, can you please take a look? Thanks. Hi @Doomski99, can you try it this way and see if it satisfies your performance requirements?

```py
import torch
from torch import nn
import ai_edge_torch


class SimpleModel(nn.Module):
    def __init__(self, input_size, hidden_size):
        super().__init__()
    
        self.lstm = nn.LSTM(input_size, hidden_size)
        self.d1 = nn.Linear(hidden_size, 1)

    def forward(self, x):
      
        x, (h0, c0) = self.lstm(x)
        x = self.d1(x)
    
        return x, torch.stack([h0, c0])


model = SimpleModel(256, 64)
sample_inputs = (torch.randn(16, 43, 256),)

edge_model = ai_edge_torch.convert(model.eval(), sample_inputs)
edge_model.export("simple_model.tflite")
```

You can find more information here: [AI-Edge-Torch](https://github.com/google-ai-edge/ai-edge-torch) and [googleblog](https://developers.googleblog.com/en/ai-edge-torch-high-performance-inference-of-pytorch-models-on-mobile-devices/) Hello @pkgoogle . Interesting library, thanks for sharing! I tested it and it seems to unroll the LSTM instead of using a loop (while) which is not quite what I was looking for. My goal was to be able to have the "UnidirectionalSequenceLSTM" Operator instead.
![image](https://github.com/tensorflow/tensorflow/assets/45564681/14a44ae3-29b8-4385-bb96-b227bb68a8e2)

### 1. System information

- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Windows 11
- TensorFlow installation (pip package or built from source): pip
- TensorFlow library (version, if pip package or github SHA, if built from source): 2.15.0

### 2. Code

```
import numpy as np
import tensorflow as tf
from tensorflow.keras import Model
from tensorflow.keras.layers import Dense, LSTM

model_name = "1x_LSTM_64_float32"
input_length = 256

class SimpleModel(Model):
  def __init__(self, input_shape, hidden_size):
    super().__init__()
    
    self.lstm = LSTM(hidden_size, return_sequences = True, return_state=True, input_shape = [-1, input_shape] )
    
    self.d1 = Dense(1, input_shape = [-1, hidden_size])

  def call(self, x):
      
    x, h0, c0 = self.lstm(x)
    x = self.d1(x)
    
    return x#, tf.stack([h0, c0])

model = SimpleModel(input_length, 64)

out, states = model(tf.random.uniform([16,43,256]))

print(np.mean(out))

model_path = f"{model_name}.tf"

run_model = tf.function(lambda x: model(x))
BATCH_SIZE = 16
STEPS = 43
INPUT_SIZE = 256
concrete_func = run_model.get_concrete_function(
    tf.TensorSpec([BATCH_SIZE, STEPS, INPUT_SIZE], tf.float32))

model.save(model_path, save_format = 'tf', signatures=concrete_func)

converter = tf.lite.TFLiteConverter.from_saved_model(model_path)
converter.target_spec.supported_ops = [
  tf.lite.OpsSet.TFLITE_BUILTINS,
]

tflite_model = converter.convert()
open(f"{model_name}.tflite", "wb").write(tflite_model)
```

The previous model produces the following graph:
![image](https://github.com/tensorflow/tensorflow/assets/45564681/4e719264-7785-41f5-8a4d-4ebb6bab990d)

With a simple hidden states retention by modifying the following code:

```
class SimpleModel(Model):
  def __init__(self, input_shape, hidden_size):
    super().__init__()
    
    self.lstm = LSTM(hidden_size, return_sequences = True, return_state=True, input_shape = [-1, input_shape] )
    
    self.d1 = Dense(1, input_shape = [-1, hidden_size])

  def call(self, x):
      
    x, h0, c0 = self.lstm(x)
    x = self.d1(x)
    
    return x, tf.stack([h0, c0])

model = SimpleModel(input_length, 64)

out, states = model(tf.random.uniform([16,43,256]))
```

I get the following graph:
![image](https://github.com/tensorflow/tensorflow/assets/45564681/d8b04b99-9bdb-4474-9b82-ecba766d2f5c)


### 3. Failure after conversion

The While OPs fail to run on my target platform unlike the UnidirectionalSequenceLstm OP.
 I can only assume that the 'UnidirectionalSequenceLstm' OP doesn't support returning the hidden states? Or am I doing something wrong? My end goal is to feed initial states to the LSTM and retain the returned ones. And I want optimized LSTM Operators. Do I need to write my own conversion logic?
 


