Apply changes from code review

Author: marcvanzee

tfjs-converter/python/tensorflowjs/BUILD

@@ -50,7 +50,10 @@ py_library(
     # This is a dummy rule used as a JAX dependency in open-source.
     # We expect JAX to already be installed on the system, e.g. via
     # `pip install jax`.
-    deps = [requirement("jax")],
+    deps = [
+        requirement("jax"),
+        requirement("importlib_resources"),
+    ],
 )
 
 py_library(

tfjs-converter/python/tensorflowjs/converters/jax_conversion.py

@@ -19,7 +19,7 @@
 from jax.experimental import jax2tf
 from jax.experimental.jax2tf import shape_poly
 import tensorflow as tf
-from tensorflowjs.converters import convert_tf_saved_model
+from tensorflowjs.converters import tf_saved_model_conversion_v2 as saved_model_conversion
 
 
 _TF_SERVING_KEY = tf.saved_model.DEFAULT_SERVING_SIGNATURE_DEF_KEY
@@ -146,4 +146,4 @@ def convert_jax(
         saved_model_dir,
         signatures=signatures,
         options=saved_model_options)
-    convert_tf_saved_model(saved_model_dir, model_dir)
+    saved_model_conversion.convert_tf_saved_model(saved_model_dir, model_dir)