docstr: Save and restore the custom gradient, when combined with other TF code.
docstr_tok: ["Save","and","restore","the","custom","gradient",",","when","combined","with","other","TF","code","."]
code: def test_gradient_nested(self):
      """Save and restore the custom gradient, when combined with other TF code."""
  
      @jax.custom_jvp
      def f_jax(x):
          return x * x
  
      @f_jax.defjvp
      def f_jax_jvp(primals, tangents):
          # 3 * x * x_t
          (x,) = primals
          (x_dot,) = tangents
          primal_out = f_jax(x)
          tangent_out = x * x_dot * 3.0
          return primal_out, tangent_out
  
      model = tf.Module()
      # After conversion, we wrap with some pure TF code
      model.f = tf.function(
          lambda x: tf.math.sin(
              jax2tf.convert(f_jax, with_gradient=True)(x)
          ),
          autograph=False,
          input_signature=[tf.TensorSpec([], tf.float32)],
      )
      f_jax_equiv = lambda x: jnp.sin(f_jax(x))
      x = np.array(0.7, dtype=jnp.float32)
      self.assertAllClose(model.f(x), f_jax_equiv(x))
      restored_model = tf_test_util.SaveAndLoadModel(model)
      xv = tf.Variable(x)
      self.assertAllClose(restored_model.f(x), f_jax_equiv(x))
      with tf.GradientTape() as tape:
          y = restored_model.f(xv)
      self.assertAllClose(
          tape.gradient(y, xv).numpy(), jax.grad(f_jax_equiv)(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 3 for calling the SavedModelTest obj's underlying member method code with example usage: obj.test_gradient_nested() and returns: unknown, primal_out,  tangent_out