code: def test_grad_pytree(self):
      call_tf = call_tf_full_ad
  
      def f_jax(xy):
          dict_ab = dict(a=2.0 * xy[0], b=xy[0] * xy[1])
          return 3.0 * dict_ab["a"] + 4.0 * dict_ab["b"]
  
      def f_outside(xy):
          dict_ab = call_tf(
              lambda xy: dict(a=2.0 * xy[0], b=xy[0] * xy[1]),
              xy,
              result_shape=dict(a=xy[0], b=xy[1]),
          )
          return 3.0 * dict_ab["a"] + 4.0 * dict_ab["b"]
  
      xy = (5.0, 6.0)
      self.assertAllClose(f_jax(xy), f_outside(xy))
      res_jax = jax.grad(f_jax)(xy)
      self.assertAllClose(res_jax, jax.grad(f_outside)(xy))
  
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 4 for calling the CallToTFTest obj's underlying member method code with example usage: obj.test_grad_pytree() and returns: unknown, unknown