code: def test_save_without_embedding_params(self):
      def model_jax(params, inputs):
          return params[0] + params[1] * inputs
  
      params = (
          np.array(1.0, dtype=jnp.float32),
          np.array(2.0, dtype=jnp.float32),
      )
      params_vars = tf.nest.map_structure(tf.Variable, params)
  
      prediction_tf = lambda x: jax2tf.convert(model_jax)(
          params_vars, x
      )
  
      model = tf.Module()
      model._variables = tf.nest.flatten(params_vars)
      model.f = tf.function(
          prediction_tf, jit_compile=True, autograph=False
      )
  
      x = np.array(0.7, dtype=jnp.float32)
      self.assertAllClose(model.f(x), model_jax(params, x))
      restored_model = tf_test_util.SaveAndLoadModel(
          model, save_gradients=False
      )
      self.assertAllClose(restored_model.f(x), model_jax(params, x))
  
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 6 for calling the SavedModelTest obj's underlying member method code with example usage: obj.test_save_without_embedding_params() and returns: unknown