code: def call_tf_full_ad(tf_fun: Callable, arg, *, result_shape):
      """Calls a TensorFlow function with support for reverse AD.
  
      Supports higher-order AD and pytree arguments.
      """
  
      @jax.custom_vjp
      def make_call(arg):
          """We wrap it all in `make_call` so that we can attach custom VJP."""
          return call_tf_no_ad(tf_fun, arg, result_shape=result_shape)
  
      # Define the fwd and bwd custom_vjp functions
      def make_call_vjp_fwd(arg):
          return (
              make_call(arg),
              arg,
          )  # Return the primal argument as the residual
  
      def make_call_vjp_bwd(res, ct_res):
          arg = res  # residual is the primal argument
  
          def tf_vjp_fun(arg_and_ct_res):
              """Invoke TF gradient; used with hcb.call."""
              arg, ct_res = arg_and_ct_res
  
              def make_var(a):
                  return (
                      a
                      if isinstance(a, tf.Variable)
                      else tf.Variable(a)
                  )
  
              arg_var = tf.nest.map_structure(make_var, arg)
  
              with tf.GradientTape(persistent=True) as tape:
                  res = tf_fun(arg_var)
  
              tf.nest.assert_same_structure(res, ct_res)
              accumulator = None  # Accumulate argument cotangent. Same structure as "arg"
  
              def acc_ct(res_, ct_res_):
                  dres_darg = tape.gradient(
                      res_,
                      sources=arg_var,
                      unconnected_gradients=tf.UnconnectedGradients.ZERO,
                  )
                  tf.nest.assert_same_structure(dres_darg, arg)
                  scaled_dres_darg = tf.nest.map_structure(
                      lambda d: d * ct_res_, dres_darg
                  )
                  nonlocal accumulator
                  accumulator = (
                      scaled_dres_darg
                      if accumulator is None
                      else tf.nest.map_structure(
                          lambda x, y: x + y,
                          accumulator,
                          scaled_dres_darg,
                      )
                  )
  
              tf.nest.map_structure(acc_ct, res, ct_res)
              return accumulator
  
          return (
              call_tf_full_ad(
                  tf_vjp_fun, (arg, ct_res), result_shape=arg
              ),
          )
  
      make_call.defvjp(make_call_vjp_fwd, make_call_vjp_bwd)
      return make_call(arg)
  
imports: from typing import Callable
  import unittest
  from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  from jax import config
  from jax import numpy
  from jax._src import test_util
  from jax._src import xla_bridge
  from jax.experimental import host_callback
  import numpy
label_desc: Use image node_id 3 for calling a global function with example usage: call_tf_full_ad(tf_fun, arg) and returns: make_call, call_tf_no_ad, make_call,  arg, call_tf_full_ad, accumulator, unknown