docstr: Conversion of function with custom VJP, inside cond.
  This exercises the custom_vjp_call_jaxpr primitives.
docstr_tok: ["Conversion","of","function","with","custom","VJP",",","inside","cond",".","This","exercises","the","custom_vjp_call_jaxpr","primitives","."]
code: def test_cond_custom_vjp(self):
      """Conversion of function with custom VJP, inside cond.
      This exercises the custom_vjp_call_jaxpr primitives."""
  
      @jax.custom_vjp
      def f(x):
          return x * x
  
      # f_fwd: a -> (b, residual)
      def f_fwd(x):
          return f(x), 3.0 * x
  
      # f_bwd: (residual, CT b) -> [CT a]
      def f_bwd(residual, ct_b):
          return (residual * ct_b,)
  
      f.defvjp(f_fwd, f_bwd)
  
      def g(x):
          return lax.cond(True, f, lambda y: y, x)
  
      arg = 0.7
      self.TransformConvertAndCompare(g, arg, None)
      self.TransformConvertAndCompare(g, arg, "vmap")
      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 6 for calling the ControlFlowOpsTest obj's underlying member method code with example usage: obj.test_cond_custom_vjp() and returns: unknown, f,  unknown, unknown, lax