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))
Tests that custom_grad can handle structured inputs/outputs.
trax/tf_numpy/extensions/extensions_test.py