docstr: Show how to combine TF-uncompilableble code with compiled JAX-converted code.
docstr_tok: ["Show","how","to","combine","TF-uncompilableble","code","with","compiled","JAX-converted","code","."]
code: def test_tf_mix_jax_with_uncompilableble(self):
      """Show how to combine TF-uncompilableble code with compiled JAX-converted code."""
  
      def tf_fn(x_str, compute_tf_fn=lambda x: x):
          # Some TF preprocessing code that cannot be compiled with XLA because it
          # uses strings.
          numbers_f32 = tf.strings.to_number(x_str, out_type=tf.float32)
          numbers_f16 = tf.cast(numbers_f32, tf.float16)
          return compute_tf_fn(numbers_f16)
  
      x_str = np.array(["3.14", "2.78"])
  
      # Test that we get an error if we try to TF-compile `tf_fn`
      with self.assertRaisesRegex(
          Exception,
          "Detected unsupported operations when trying to compile graph",
      ):
          tf.function(tf_fn, jit_compile=True, autograph=False)(x_str)
  
      # Plug in the TF-compiled JAX-converted `compute_jax_fn`.
      composed_fn = lambda x_str: tf_fn(
          x_str,
          compute_tf_fn=tf.function(
              jax2tf.convert(jnp.sin), autograph=False, jit_compile=True
          ),
      )
      res_tf = composed_fn(x_str)
      self.assertAllClose(
          res_tf.numpy(),
          jnp.sin(np.array([3.14, 2.78], dtype=np.float16)),
      )
  
      # Save and restore SavedModel
      restored_f, _ = tf_test_util.SaveAndLoadFunction(
          composed_fn, input_args=[x_str]
      )
      res_tf_restored = restored_f(x_str)
      self.assertAllClose(res_tf_restored.numpy(), res_tf.numpy())
  
imports: from absl.testing import absltest
  import os
  import jax
  from jax import lax
  import jax.numpy
  import numpy
  import tensorflow
  from jax.experimental import jax2tf
  from jax.experimental.jax2tf.tests import tf_test_util
  from jax._src import test_util
  from jax import config
label_desc: Use image node_id 12 for calling the SavedModelTest obj's underlying member method code with example usage: obj.test_tf_mix_jax_with_uncompilableble() and returns: compute_tf_fn