Bug: Universal Sentence Encoder breaks when using Tensorflow 2.10.0

[ProfDoof]

What happened?

System Info: Web App: Google CoLab Jupyter Runtime: Local Jupyter Notebook running on Python 3.10 Tensorflow Version: 2.10.0

I'm trying to use the model from tfhub.dev/google/universal-sentence-encoder/4 for embedding sentences to use for multiple downstream tasks. While trying to do the embedding process with my local runtime, I ran into a JIT compilation error.

It is as follows

JIT compilation failed.
     [[{{node EncoderDNN/EmbeddingLookup/EmbeddingLookupUnique/embedding_lookup/mod}}]] [Op:__inference_restored_function_body_13685]

I expected this to just work but I may have a misunderstanding of how the models are expected to be used.

Relevant code

import tensorflow_hub as hub
# Load pre-trained universal sentence encoder model
embed = hub.load("https://tfhub.dev/google/universal-sentence-encoder/4")
embed(['test'])

Relevant log output

2022-11-15 09:58:37.233443: W tensorflow/core/framework/op_kernel.cc:1768] UNKNOWN: JIT compilation failed.
---------------------------------------------------------------------------
UnknownError                              Traceback (most recent call last)
Cell In [27], line 3
      1 # Load pre-trained universal sentence encoder model
      2 embed = hub.load("https://tfhub.dev/google/universal-sentence-encoder/4")
----> 3 embed(['test'])

File ~/.local/lib/python3.10/site-packages/tensorflow/python/saved_model/load.py:704, in _call_attribute(instance, *args, **kwargs)
    703 def _call_attribute(instance, *args, **kwargs):
--> 704   return instance.__call__(*args, **kwargs)

File ~/.local/lib/python3.10/site-packages/tensorflow/python/util/traceback_utils.py:153, in filter_traceback.<locals>.error_handler(*args, **kwargs)
    151 except Exception as e:
    152   filtered_tb = _process_traceback_frames(e.__traceback__)
--> 153   raise e.with_traceback(filtered_tb) from None
    154 finally:
    155   del filtered_tb

File ~/.local/lib/python3.10/site-packages/tensorflow/python/eager/execute.py:54, in quick_execute(op_name, num_outputs, inputs, attrs, ctx, name)
     52 try:
     53   ctx.ensure_initialized()
---> 54   tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
     55                                       inputs, attrs, num_outputs)
     56 except core._NotOkStatusException as e:
     57   if name is not None:

UnknownError: Graph execution error:

JIT compilation failed.
     [[{{node EncoderDNN/EmbeddingLookup/EmbeddingLookupUnique/embedding_lookup/mod}}]] [Op:__inference_restored_function_body_13685]

tensorflow_hub Version

0.12.0 (latest stable release)

TensorFlow Version

other (please specify)

Other libraries

tensorflow==2.10.0

Full pip freeze dump can be found here pip freeze dump on hatebin

Python Version

3.x

OS

Linux

[singhniraj08]

@ProfDoof,

By default colab uses python 3.7, So I installed python 3.10 with tensorflow 2.10 and tensorflow_hub 0.12.0 and executed the provided snippet.

I was unable to reproduce the error mentioned and got the encoded vector from USE successfully. Please find attach gist.

Could you please try reinstalling the tensorflow build and run the same code to see if the issue persists?

Thank you!

[ProfDoof]

I have already tried forcing a reinstall of the tensorflow build. Is it possible that it's caused by my having a GPU or something?

[raiyaan555]

https://discuss.tensorflow.org/t/how-to-solve-tensorflow-jit-compile-error/3707

Hey maybe this can be of help to you

[ProfDoof]

I saw that but it's not my model so I'm not sure I can do anything about it.

[raiyaan555]

I have already tried forcing a reinstall of the tensorflow build. Is it possible that it's caused by my having a GPU or something?

GPU should not be a problem as such.

[ProfDoof]

Lol, well then I don't know what the issue is. I will try it again as soon as I get home. I'm just getting done with a dentist appointment 😂

[ProfDoof]

Figured it out. For some godforsaken reason, I was setting CUDA_DIR to the correct location but I kept getting this error "error: Can't find libdevice directory ${CUDA_DIR}/nvvm/libdevice" which turns out is because I had to do "export XLA_FLAGS=--xla_gpu_cuda_data_dir=/usr/lib/cuda-11.2" instead.

Now, I'm not sure why this is the case and it would be greatly appreciated if someone could point me in the direction of the right people to bug with this bug to get it fixed but that's the close of this bug here for me.