docstr: Tests that custom_grad can handle structured inputs/outputs.
docstr_tok: ["Tests","that","custom_grad","can","handle","structured","inputs\/outputs","."]
code: def testCustomGradStructure(self, decorator, x_struct, y_struct):
      """Tests that custom_grad can handle structured inputs/outputs."""
  
      def zeros(x):
          return tf.nest.map_structure(
              lambda _: tf_np.zeros([], np.float32), x
          )
  
      def get_struct(x):
          return tf.nest.map_structure(lambda _: None, x)
  
      @extensions.custom_grad
      def f(*x):
          del x
  
          def vjp(dy):
              self.assertEqual(y_struct, get_struct(dy))
              return zeros(x_struct)
  
          return zeros(y_struct), vjp
  
      x, dy = zeros([x_struct, y_struct])
  
      @decorator
      def run(x, dy):
          y, vjp = extensions.vjp(f, *x)
          dx = vjp(dy)
          return dx, y
  
      dx, y = run(x, dy)
      self.assertEqual(x_struct, get_struct(dx))
      self.assertEqual(y_struct, get_struct(y))
  
imports: import functools
  import itertools
  from absl import flags
  from absl.testing import parameterized
  from jax import lax
  import numpy
  import tensorflow.compat.v2
  from trax.tf_numpy import extensions
  import trax.tf_numpy.numpy
label_desc: Use image node_id 4 for calling the ExtensionsTest obj's underlying member method code with example usage: obj.testCustomGradStructure(decorator, x_struct, y_struct) and returns: tf, tf, zeros,  vjp, dx,  y, zeros