TENSORFLOWJS CONVERTER NOT WORK IN GOOGLE COLAB

[JesusPe1917]

Hello, good afternoon, I open this space to report an error in tensorflowjs converter, it appears in Google Colab when I try to convert my .h5 model to tensorflowjs, the error is shown below

2023-12-24 17:30:00.401498: E external/local_xla/xla/stream_executor/cuda/cuda_dnn.cc:9261] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2023-12-24 17:30:00.401551: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:607] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2023-12-24 17:30:00.403137: E external/local_xla/xla/stream_executor/cuda/cuda_blas.cc:1515] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2023-12-24 17:30:01.846435: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
Traceback (most recent call last):
File "/usr/local/bin/tensorflowjs_converter", line 5, in
from tensorflowjs.converters.converter import pip_main
File "/usr/local/lib/python3.10/dist-packages/tensorflowjs/init.py", line 21, in
from tensorflowjs import converters
File "/usr/local/lib/python3.10/dist-packages/tensorflowjs/converters/init.py", line 26, in
from tensorflowjs.converters.jax_conversion import convert_jax
File "/usr/local/lib/python3.10/dist-packages/tensorflowjs/converters/jax_conversion.py", line 20, in
from jax.experimental.jax2tf import shape_poly
ImportError: cannot import name 'shape_poly' from 'jax.experimental.jax2tf' (/usr/local/lib/python3.10/dist-packages/jax/experimental/jax2tf/init.py)

I have already tried from another notebook, or by reinstalling the library, the code worked without problems but a few days ago it stopped working, I suspect it is due to the python update in Google Colab a few days ago

[gaikwadrahul8]

Hi, @JesusPe1917

Thank you for bringing this issue to our attention, Could you please help me with your model in zip format so I'll try to replicate the same behavior from my end to investigate this issue further as soon as possible. Thank you for your understanding and patience.

[JesusPe1917]

Hi! Gaikwadrahul8. Thank you very much for taking care of my case, of course! I attach the zip of the model
modelo.zip

I apologize in advance because I don't know how to use this platform very well, So I'm not sure I'm answering you correctly.

Hola,@JesusPe1917

Gracias por informarnos sobre este problema. ¿Podría ayudarme con su modelo en formato zip para intentar replicar el mismo comportamiento por mi parte para investigar este problema más a fondo lo antes posible? Gracias por su comprensión y paciencia.

[andreariba]

I have the same issue, I believe it depends on this issue jax-ml/jax#18978 from jax.

[andreariba]

just tested if you downgrade jax and jaxlib to 0.4.21 it works, so temporarily ...

[JesusPe1917]

Wow!, amazing!, sorry, I'm still new to all this, could you tell me how you downgraded jax? I guess that's what can be done until they fix it.

[gaikwadrahul8]

Hi, @andreariba Thank you for so much for your insights and workaround, I really appreciate your efforts

Hi, @JesusPe1917 I have created gist-file for your reference and things are working as expected after downgrading the jax and jaxlib to 0.4.21 and model is converting successfully to tfjs layers-model format so please give it try from your end and please let me know, Is it working as expected or not with that temporary workaround downgrading the jax and jaxlib to 0.4.21 ?

At the moment things are not working with jax and jaxlib to 0.4.23 so our relevant team will work on this issue and will fix this issue as soon as possible. Thank you for bringing this issue to our attention

Thank you for your understanding and patience.

[gaikwadrahul8]

Hi, @JesusPe1917

Could you please confirm tensorflowjs_converter working as expected after downgrading the jax and jaxlib to 0.4.21 ? If so please feel free to close this issue

If you're facing any issue please let us know with error log to help you further. Thank you for your understanding and patience.