docstr: Conversion of function with custom JVP, inside cond.
  This exercises the custom_jvp_call_jaxpr primitives.
docstr_tok: ["Conversion","of","function","with","custom","JVP",",","inside","cond",".","This","exercises","the","custom_jvp_call_jaxpr","primitives","."]
code: def test_cond_custom_jvp(self):
      """Conversion of function with custom JVP, inside cond.
      This exercises the custom_jvp_call_jaxpr primitives."""
  
      @jax.custom_jvp
      def f(x):
          return x * x
  
      @f.defjvp
      def f_jvp(primals, tangents):
          (x,) = primals
          (x_dot,) = tangents
          primal_out = f(x)
          tangent_out = 3.0 * x * x_dot
          return primal_out, tangent_out
  
      def g(x):
          return lax.cond(True, f, lambda y: y, x)
  
      arg = 0.7
      self.TransformConvertAndCompare(g, arg, None)
      self.TransformConvertAndCompare(g, arg, "jvp")
      self.TransformConvertAndCompare(g, arg, "vmap")
      self.TransformConvertAndCompare(g, arg, "jvp_vmap")
      self.TransformConvertAndCompare(g, arg, "grad")
      self.TransformConvertAndCompare(g, arg, "grad_vmap")
  
imports: from absl.testing import absltest
  import jax
  import jax.lax
  import jax.numpy
  from jax._src import test_util
  import numpy
  from jax.experimental.jax2tf.tests import tf_test_util
  from jax import config
label_desc: Use image node_id 5 for calling the ControlFlowOpsTest obj's underlying member method code with example usage: obj.test_cond_custom_jvp() and returns: unknown, primal_out,  tangent_out, lax