code: def call_tf_simple_ad(tf_fun: Callable, arg, *, result_shape):
      """Calls a TensorFlow function with simple support for reverse AD.
  
      Works only for 1st order AD and only for arguments and results being a single
      ndarray (no pytrees). Functions whose name starts with "tf_" are TensorFlow
      functions and must be called outside the JAX computation.
      """
  
      @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 the primal argument as the residual. Use `make_call` for the
          # primal computation to enable higher-order AD.
          return make_call(arg), arg
  
      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
              arg_var = tf.Variable(arg)
              with tf.GradientTape(persistent=True) as tape:
                  res = tf_fun(arg_var)
  
              dres_darg = tape.gradient(
                  res,
                  sources=arg_var,
                  output_gradients=ct_res,
                  unconnected_gradients=tf.UnconnectedGradients.ZERO,
              )
              return dres_darg
  
          return (
              call_tf_simple_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 2 for calling a global function with example usage: call_tf_simple_ad(tf_fun, arg) and returns: make_call, call_tf_no_ad, make_call,  arg, call_tf_simple_ad, dres_darg