ValueError when taking the gradient of a jax2tf-converted function that takes a pytree_node_class argument

[magnusdk]

Hi!

I am experiencing an issue with taking the gradient of a Tensorflow-converted function using tensorflow.GradientTape when one of the function arguments is a custom pytree_node_class. There's a mismatch with the number of outputs and the number of generated gradients, presumably because the custom pytree_node_class is flattened into multiple parameters.

Below is a minimal(-ish) reproducable example:

import jax.numpy as jnp
import tensorflow as tf
from jax.experimental import jax2tf
from jax.tree_util import register_pytree_node_class


@register_pytree_node_class
class Special(tf.Module):
    def __init__(self, x, y):
        self.x = x
        self.y = y
    def tree_flatten(self):
        return ((self.x, self.y), None)
    @classmethod
    def tree_unflatten(cls, aux_data, children):
        return cls(*children)

    def do_thing(self, a: jnp.ndarray):
        return self.x * a + self.y


@jax2tf.convert
def my_function(my_obj: Special, a: jnp.ndarray):
    return jnp.sum(my_obj.do_thing(a))


my_obj = Special(tf.Variable(2.0), tf.Variable(3.0))
with tf.GradientTape() as tape:
    # This next line raises a ValueError
    result = my_function(my_obj, tf.range(10))
grads = tape.gradient(result, [my_obj.trainable_variables])

Result:

Traceback (most recent call last):
  File ".../minimal_example.py", line 34, in <module>
    grads = tape.gradient(result, [my_obj.trainable_variables])
  File ".../python3.8/site-packages/tensorflow/python/eager/backprop.py", line 1084, in gradient
    flat_grad = imperative_grad.imperative_grad(
  File ".../python3.8/site-packages/tensorflow/python/eager/imperative_grad.py", line 71, in imperative_grad
    return pywrap_tfe.TFE_Py_TapeGradient(
  File ".../python3.8/site-packages/tensorflow/python/ops/custom_gradient.py", line 572, in actual_grad_fn
    raise ValueError(
ValueError: ('custom_gradient function expected to return', 3, 'gradients but returned', 2, 'instead.')

I really like the ability to register custom classes for JAX for improving code readability and maintainability. I am working on a project that requires interop with Tensorflow, however.

Is there a way to make this work while keeping my_function's call signature?