code: def test_save_without_gradients(self):
      f_jax = lambda x: x * x
  
      x = np.array(0.7, dtype=jnp.float32)
      model = tf.Module()
      model.f = tf.function(
          jax2tf.convert(f_jax, with_gradient=True),
          autograph=False,
          input_signature=[tf.TensorSpec(x.shape, x.dtype)],
      )
  
      self.assertAllClose(model.f(x), f_jax(x))
      restored_model = tf_test_util.SaveAndLoadModel(
          model, save_gradients=False
      )
      self.assertAllClose(restored_model.f(x), f_jax(x))
  
      xv = tf.Variable(x)
      with tf.GradientTape():
          _ = restored_model.f(xv)
  
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 5 for calling the SavedModelTest obj's underlying member method code with example usage: obj.test_save_without_gradients() without return types