Super Slow Performance (tf.function fails) of GradientTape with LSTM and Manual Training Loop
@kasrash 
I tried to replicate the issue on Colab using TF v2.11 but facing a different problem. Please find the gist [here](https://colab.research.google.com/gist/tiruk007/b415002491a70d80b4181d00aeef3369/untitled123.ipynb) for reference.
Could you please provide all dependencies to replicate the issue reported here ?

Thank you! > @kasrash I tried to replicate the issue on Colab using TF v2.11 but facing a different problem. Please find the gist [here](https://colab.research.google.com/gist/tiruk007/b415002491a70d80b4181d00aeef3369/untitled123.ipynb) for reference. Could you please provide all dependencies to replicate the issue reported here ?
> 
> Thank you!

Hi Tiruk007,

Thank you for checking this. I created a gist at [here](https://colab.research.google.com/gist/kasrash/5ad0b2e8df48b288c82d98be7b59803b/tferrorreplicate.ipynb), and you can download the required data [here](https://drive.google.com/drive/folders/1IG5q0U1xCqVjpC_35F4j0uER5rackoNK?usp=share_link). However, I had to reduce the model size so that it won't run out of Ram. As you can see the issue is not created in the Colab. I am not sure if some of the outputs are suppressed there or there is an issue with my TF installation. I would be grateful to know your thoughts.
 @kasrash 
Sorry for the late reply. I tried to execute the given code  on a Linux machine using 2.11 where I got 6.47 seconds for each step. Please find the [screenshot](https://user-images.githubusercontent.com/111861663/219692401-c6e1b83e-b0c4-4a9c-a845-7af47412d5f0.png) for reference.Could you please make sure to check the tested build configuration as mentioned [here](https://www.tensorflow.org/install/source#tested_build_configurations) and use the latest TF version.

Thank you !

 > @kasrash Sorry for the late reply. I tried to execute the given code on a Linux machine using 2.11 where I got 6.47 seconds for each step. Please find the [screenshot](https://user-images.githubusercontent.com/111861663/219692401-c6e1b83e-b0c4-4a9c-a845-7af47412d5f0.png) for reference.Could you please make sure to check the tested build configuration as mentioned [here](https://www.tensorflow.org/install/source#tested_build_configurations) and use the latest TF version.
> 
> Thank you !

@tiruk007 hi

Thank you so much for testing this. I see that your not getting the same error as I am getting. I will check my installation, but I installed TF 2.11 using pip and the installation guide on the website. Nevertheless, don't you thin that 6.47 s for this toy example is too high? Since this is almost the same logic as model.fit() and that takes around ms to complete each step. Is there any way to speed up this process? Hello,
Just an update on this, I tried the code and a MNIST toy model on several machines (Windows and Ubuntu) with different versions of TF (2.9 and 2.11) all give the same error.

Changing the activation from tanh to sigmoid or relu solves the issue, and the performance is blazing fast. Of course CuDNN will not be utilized.

Hope it helps in solving this issue and I would be grateful to know your inputs. Also, [here](https://colab.research.google.com/gist/kasrash/ba04bd4f82e3de6423d67f828083f03b/untitled2.ipynb) is a gist of toy MNIST model. If you go to Runtime->Runtime log you can see the error. Hi, @kasrash 

I have tried to execute the same code from gist-file of toy MNIST model with `tanh` activation function and it's working as expected here is [gist-file](https://colab.research.google.com/gist/gaikwadrahul8/09be3a67beeb4957871cefb232efb964/-59677-test.ipynb) and even I checked the runtime error log in that all I'm getting warning messages as far I know we can ignore those warnings messages unless and until the code is running as expected and giving the correct result. You can refer [stack overflow answers ](https://stackoverflow.com/questions/35911252/disable-tensorflow-debugging-information)which also supports 

If you're looking to disable `absl INFO and WARNING log messages `please use below code snippet:

```
import absl.logging
absl.logging.set_verbosity(absl.logging.ERROR)

```
Here is screenshot of Google Colab runtime error log :

![image](https://user-images.githubusercontent.com/115997457/220887073-3eb045e8-c95e-4fb8-a0c0-9517c9ab3abb.png)


If I've missed something here please let me know ? Just to confirm, May I know your initial issue slow performance with tf.fuction got resolved or not or do you still need help with that issue please ? Thank you!
 Hi @gaikwadrahul8,

That is true that it is a warning, but when I searched for this message on Google, I found several bug reports here dating back to 2019 and 2020, which TF team confirmed that it was actually a bug. Also, with "tanh" activation on my main code, the TF GradiantTape Jacobian (decorated with tf.function) calculation takes ~13 s, but when changing it to relu or sigmoid activations, it only takes ~0.9 s when decorated with tf.function and ~15 s without tf.function. So, I am guessing definitely something (maybe odd )is happening when tanh activation is used as it should be faster since it use CuDNN kernel.

I would be grateful to know your thoughts on this, and many thanks for your time. Hi, @kasrash 

Thank you for your detailed observations and if possible could you please help us with issues number which we considered warnings as bug back to 2019 and 2020 ?

If you're using `tf.function`,the behavior of `tf.function `when `AutoGraph `is used by default `AutoGraph==True `and `@tf.function` decorated function behaves as we expect if we use a `tf.Tensor `as input type (please use `tf.Tensor `everywhere!) and if the input is a Python native type creates a new graph. You can refer this [stack overflow answer](https://stackoverflow.com/questions/55711115/tensorflow-2-0-model-using-tf-function-very-slow-and-is-recompiling-every-time-t) and also refer these useful article links [Ref-1](https://pgaleone.eu/tensorflow/tf.function/2019/04/03/dissecting-tf-function-part-2/#using-a-python-native-type), [Ref-2](https://pgaleone.eu/tensorflow/tf.function/2019/05/10/dissecting-tf-function-part-3/), I hope it will help you to resolve your issue. Thank you!

 > Hi, @kasrash
> 
> Thank you for your detailed observations and if possible could you please help us with issues number which we considered warnings as bug back to 2019 and 2020 ?
> 
> If you're using `tf.function`,the behavior of `tf.function `when `AutoGraph `is used by default `AutoGraph==True `and `@tf.function` decorated function behaves as we expect if we use a `tf.Tensor `as input type (please use `tf.Tensor `everywhere!) and if the input is a Python native type creates a new graph. You can refer this [stack overflow answer](https://stackoverflow.com/questions/55711115/tensorflow-2-0-model-using-tf-function-very-slow-and-is-recompiling-every-time-t) and also refer these useful article links [Ref-1](https://pgaleone.eu/tensorflow/tf.function/2019/04/03/dissecting-tf-function-part-2/#using-a-python-native-type), [Ref-2](https://pgaleone.eu/tensorflow/tf.function/2019/05/10/dissecting-tf-function-part-3/), I hope it will help you to resolve your issue. Thank you!

Hi @gaikwadrahul8 

Please see issues [#35928](https://github.com/tensorflow/tensorflow/issues/35928), [#41129](https://github.com/tensorflow/tensorflow/issues/41129), and [#41949](https://github.com/tensorflow/tensorflow/issues/41949). There was also a PR ([here](https://github.com/tensorflow/tensorflow/commit/3a03164a24ea7dbfe221e84ca62e2a39b745c379), which not sure if it is relevant.

Thank you for the suggestion. But, the input to the function decorated with @tf.function is a Tensor. Also, it seems odd to me that this is happening only when using tanh activation (hence, CuDNNLSTM), but not with other activations (i.e., not CuDNNLSTMs). Hi, @kasrash

Thank you for sharing those similar issues, if possible could you please help me with your complete code snippet in Google Colab notebook or as gist-file with all imports so I'll test it from my end because I tested above toy `MNIST` model  it's working fine with `tanh` activation functions ? also please help me with Tensorlflow, Python versions ? I hope you're using  `cudatoolkit=11.2` and `cudnn=8.1.0 `which is tested configurations as per our [official documentation ](https://www.tensorflow.org/install/source#gpu).  

You can also refer our official documentation for [Speeding-up your training step with tf.function](https://www.tensorflow.org/guide/keras/writing_a_training_loop_from_scratch#speeding-up_your_training_step_with_tffunction) and You can also try [tf.function(jit_compile=True)](https://www.tensorflow.org/api_docs/python/tf/function) for a more significant performance boost, especially if your code is heavy on TensorFlow control flow and uses many small tensors. Learn more in the _Explicit compilation with [tf.function(jitcompile=True)](https://www.tensorflow.org/api_docs/python/tf/function) section of the [XLA overview](https://www.tensorflow.org/xla). Thank you!

 This issue was closed because it has been inactive for 7 days since being marked as stale. Please reopen if you'd like to work on this further. Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/59677">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/59677">No</a>

<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Bug

### Have you reproduced the bug with TF nightly?

No

### Source

binary

### Tensorflow Version

TF 2.11.0

### Custom Code

No

### OS Platform and Distribution

Linux Ubuntu 20.04

### Mobile device

_No response_

### Python version

3.9.16

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

12.0 (V12.0.140)

### GPU model and memory

2* Nvidia RTX A6000 48 GB

### Current Behaviour?

```shell
Hello all, 
For a specific problem, I need a manual training loop that includes a intermediate function. This function gets the model jacobian (output w.r.t to model's weights) along with other variables to generate the gradients and the loss function. I tried to calculate the model's Jacobian with GradientTape inside a function wrapped by tf.function. However, I get the error that tf.function fails. I saw similar issues dating back to 2020, but they were solved back then. I would be grateful to know your thoughts as the performance of the current code is super slow (~13 seconds for each step).
P.S. I cannot wrap the whole train_step function below with tf.function as the intermediate functions use Python dicts and Numpy arrays (the algorithm is a recursive algorithm if it matters).
Many thanks in advance!
```


### Standalone code to reproduce the issue

```shell
#fn to calculate the model jacobian
@tf.function
def jacobian_calc(model, inputs):
    with tf.GradientTape() as tape:
        tape.watch(model.trainable_weights)
        predictions = model(inputs, training=True)
    model_jacobian = tape.jacobian(predictions,model.trainable_weights)
    return model_jacobian, predictions

def train_step(model, inputs, targets, acc):
    st = time.time()
    model_jacobian, predictions = jacobian_calc(model,inputs)
    et = time.time()
    print(f"Elapsed time is {et-st}")
    dLoss = []
    for jacobian in model_jacobian:
        p_loss, p_dLoss = loss_grads(targets, predictions, acc, jacobian)
        p_dLoss = tf.convert_to_tensor(p_dLoss, dtype=tf.float32)
        print(type(p_dLoss))
        print(type(model.trainable_weights[0]))
        dLoss.append(p_dLoss/K.eval(inputs).size)
    Loss = p_loss/K.eval(inputs).size
    opt.apply_gradients(zip(dLoss, model.trainable_weights))
    logs = {}
    print(f" Loosssss is = {Loss}")
    logs["loss"] = Loss
    return logs

# load the samples
with open(seq_path+'/real_1000_300.pkl','rb') as dumper:
    samples_real = pickle.load(dumper)
with open(seq_path+'/model_1000_300.pkl','rb') as dumper:
    samples_model = pickle.load(dumper)
with open(seq_path+'/acc_1000_300.pkl','rb') as dumper:
    samples_acc = pickle.load(dumper)

# create training and testing sets
train_x = samples_model[:val_index_opt, :, 0]    #Umodel
train_y = samples_real[:val_index_opt, :, 0]     #Ureal
train_acc = samples_acc[:val_index_opt, :, 0]    #acc
val_x = samples_model[val_index_opt:, :, 0]     #Umodel
val_y = samples_real[val_index_opt:, :, 0]      #Ureal
val_acc = samples_acc[val_index_opt:, :, 0]     #acc

# no data scaling as I need to figure this out later. A dirty fix for now! Sorry:)
train_scaled_x = train_x.reshape(train_x.shape[0],train_x.shape[1],-1)
train_scaled_y = train_y.reshape(train_y.shape[0],train_y.shape[1],-1)
train_scaled_acc = train_acc.reshape(train_acc.shape[0],train_acc.shape[1],-1)
val_scaled_x = val_x.reshape(val_x.shape[0],val_x.shape[1],-1)
val_scaled_y = val_y.reshape(val_y.shape[0],val_y.shape[1],-1)
val_scaled_acc = val_acc.reshape(val_acc.shape[0],val_acc.shape[1],-1)

## Creating the RNN
model = Sequential()
                        
model.add(LSTM(32, input_shape=(None, train_scaled_x.shape[2]), mreturn_sequences=True))
                        model.add(Dense(32, input_shape=(None, train_scaled_x.shape[2])))
model.add(Dense(dense_1, activation='relu'))
model.add(Dense(1, activation='linear'))

# optimizer
opt = tf.keras.optimizers.Adam(learning_rate=0.001)

#Train model (only presented the training and not the validation here
training_set = tf.data.Dataset.from_tensor_slices((train_scaled_x,train_scaled_acc,train_scaled_y))
training_set = training_set.batch(batch)
history = {"loss": []}
for ee in range(1,epoch+1):
    for inputs_batch, acc_batch, targets_batch in training_set:
        logs = train_step(model, inputs_batch,targets_batch,acc_batch)
    print(f"Results at the end of epoch {epoch}")
    for key, value in logs.items(): #have some other metrics that I omitted here
        print(f"...{key}: {value:.4f}")
    history["loss"].append(logs["loss"])
```


### Relevant log output

```shell
2023-02-13 21:51:20.518488: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2023-02-13 21:51:21.105332: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /opt/spack/v0.17.0/opt/spack/linux-ubuntu20.04-zen2/gcc-9.3.0/slurm-21-08-1-1-ctwolps4xy7mz7h3hooji4t72xs5vsyz/lib:/home/kshamsaei/miniconda3/envs/tf/lib/
2023-02-13 21:51:21.105406: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory; LD_LIBRARY_PATH: /opt/spack/v0.17.0/opt/spack/linux-ubuntu20.04-zen2/gcc-9.3.0/slurm-21-08-1-1-ctwolps4xy7mz7h3hooji4t72xs5vsyz/lib:/home/kshamsaei/miniconda3/envs/tf/lib/
2023-02-13 21:51:21.105417: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.
HELLO
/home/kshamsaei/sdof/sdof-300-1000
2023-02-13 21:51:21.662319: I tensorflow/core/platform/cpu_feature_guard.cc:193] This TensorFlow binary is optimized with oneAPI Deep Neural Network Library (oneDNN) to use the following CPU instructions in performance-critical operations:  AVX2 FMA
To enable them in other operations, rebuild TensorFlow with the appropriate compiler flags.
2023-02-13 21:51:22.347877: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1613] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 46671 MB memory:  -> device: 0, name: NVIDIA RTX A6000, pci bus id: 0000:21:00.0, compute capability: 8.6
2023-02-13 21:51:22.348397: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1613] Created device /job:localhost/replica:0/task:0/device:GPU:1 with 46685 MB memory:  -> device: 1, name: NVIDIA RTX A6000, pci bus id: 0000:81:00.0, compute capability: 8.6
WARNING:tensorflow:From /home/kshamsaei/miniconda3/envs/tf/lib/python3.9/site-packages/tensorflow/python/autograph/pyct/static_analysis/liveness.py:83: Analyzer.lamba_check (from tensorflow.python.autograph.pyct.static_analysis.liveness) is deprecated and will be removed after 2023-09-23.
Instructions for updating:
Lambda fuctions will be no more assumed to be used in the statement where they are used, or at least in the same block. https://github.com/tensorflow/tensorflow/issues/56089
2023-02-13 21:51:25.417706: E tensorflow/core/grappler/optimizers/meta_optimizer.cc:954] function_optimizer failed: INVALID_ARGUMENT: Input 0 of node loop_body/PartitionedCall/pfor/PartitionedCall/gradients/while_grad/while_grad/pfor/PartitionedCall/while/enter/_746 was passed int32 from sequential/lstm/PartitionedCall:20 incompatible with expected variant.
2023-02-13 21:51:25.926875: E tensorflow/core/grappler/optimizers/meta_optimizer.cc:954] function_optimizer failed: INVALID_ARGUMENT: Input 0 of node loop_body/PartitionedCall/pfor/PartitionedCall/gradients/while_grad/while_grad/pfor/PartitionedCall/while/enter/_746 was passed int32 from sequential/lstm/PartitionedCall:20 incompatible with expected variant.
2023-02-13 21:51:26.001961: W tensorflow/core/common_runtime/process_function_library_runtime.cc:915] Ignoring multi-device function optimization failure: INVALID_ARGUMENT: Input 0 of node loop_body/PartitionedCall/pfor/PartitionedCall/gradients/while_grad/while_grad/pfor/PartitionedCall/while/enter/_728 was passed float from sequential/lstm/PartitionedCall:11 incompatible with expected int32.
2023-02-13 21:51:31.457799: I tensorflow/compiler/xla/stream_executor/cuda/cuda_blas.cc:630] TensorFloat-32 will be used for the matrix multiplication. This will only be logged once.
```
</details>