JVP using tf.autodiff.ForwardAccumulator becomes None under graph execution
When using **tf.autodiff.ForwardAccumulator** with **@tf.function**, the forward-pass is executed during tracing, which means that the **ForwardAccumulator** instance is not being executed during the graph execution. To fix this, you can wrap the **ForwardAccumulator** instance with **tf.custom_gradient**.
```
@tf.function
def train(self, data):
    with tf.GradientTape() as upper_tape:
        loss1= self.loss1(data)
    grad1 = upper_tape.gradient(loss1, self.net1.variables)

    def jvp_with_acc(tangents):
        with tf.autodiff.ForwardAccumulator(primals=self.net1.variables, tangents=tangents) as acc:
            with tf.GradientTape() as lower_tape:
                loss2 = self.loss2(data)
            grad2 = lower_tape.gradient(loss2, self.net2.variables)

        final_grad = acc.jvp(grad2)
        return final_grad, None

    final_grad = jvp_with_acc(grad1)
    self.optimizer.apply_gradients(zip(final_grad, self.net2.variables))
```
 > When using **tf.autodiff.ForwardAccumulator** with **@tf.function**, the forward-pass is executed during tracing, which means that the **ForwardAccumulator** instance is not being executed during the graph execution. To fix this, you can wrap the **ForwardAccumulator** instance with **tf.custom_gradient**.
> 
> ```
> @tf.function
> def train(self, data):
>     with tf.GradientTape() as upper_tape:
>         loss1= self.loss1(data)
>     grad1 = upper_tape.gradient(loss1, self.net1.variables)
> 
>     def jvp_with_acc(tangents):
>         with tf.autodiff.ForwardAccumulator(primals=self.net1.variables, tangents=tangents) as acc:
>             with tf.GradientTape() as lower_tape:
>                 loss2 = self.loss2(data)
>             grad2 = lower_tape.gradient(loss2, self.net2.variables)
> 
>         final_grad = acc.jvp(grad2)
>         return final_grad, None
> 
>     final_grad = jvp_with_acc(grad1)
>     self.optimizer.apply_gradients(zip(final_grad, self.net2.variables))
> ```

Thank you so much for replying. I managed to get the code to work under eager mode using the following code, yet when I turn on graph mode using @tf.function I got "ValueError: Tried to convert 'input' to a tensor and failed. Error: None values not supported.". I guess it is because "grad1" is None when first building the graph. Is there something I can do?

```
def train(self, data):
    with tf.GradientTape() as upper_tape:
        loss1= self.loss1(data)
    grad1 = upper_tape.gradient(loss1, self.net1.variables)
    
    @tf.custom_gradient
    def jvp_with_acc(tangents):
        with tf.autodiff.ForwardAccumulator(primals=self.net1.variables, tangents=tangents) as acc:
            with tf.GradientTape() as lower_tape:
                loss2 = self.loss2(data)
            grad2 = lower_tape.gradient(loss2, self.net2.variables)

        final_grad = acc.jvp(grad2)
        def grad(variables):
            return None
        return final_grad, grad

    final_grad = jvp_with_acc(grad1)
    self.optimizer.apply_gradients(zip(final_grad, self.net2.variables))

```  you can try explicitly converting your input to a **tf.Tensor** before computing the gradients. You can also try setting a default value for **grad1** in case it is not being computed properly > 

Turns out the error comes from somewhere else in my loss function. The code is all good now. Much thanks. Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/60164">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/60164">No</a>

I'm new to Tensorflow. I'm trying to compute a Jacobian-Vector Product using tf.autodiff.ForwardAccumulator with a train function looks something like the code below. The jvp looks fine under eager execution. However, the jvp becomes a list of Nones when activate graph execution using @tf.function. 
What could be the issue here?

### Tensorflow Version

tf 2.6.0

### Python version

3.7

### Train code

```shell
@tf.function
def train(self, data):
        with tf.GradientTape() as upper_tape:
            loss1= self.loss1(data)
        grad1 = upper_tape.gradient(loss1, self.net1.variables)

        with tf.autodiff.ForwardAccumulator(primals=self.net1.variables, tangents=grad1) as acc:
            with tf.GradientTape() as lower_tape:
                loss2 = self.loss2(data)
            grad2 = lower_tape.gradient(loss2, self.net2.variables)

        final_grad = acc.jvp(grad2)
        self.optimizer.apply_gradients(zip(final_grad, self.net2.variables))
```

grad1 and grad2 are correctly computed under both eager mode and graph mode. The only problem is with the jvp.