Forward AD threw error in tf.autodiff.ForwardAccumulator for tf.keras.layers.AveragePooling3D, but backward AD succeeded with same input
@cheyennee,
Looks like this is duplicate of issue [#56833](https://github.com/tensorflow/tensorflow/issues/56833). Could you please close this issue, since it is already being tracked there? Thank you! Hi, @tilakrayal ,
It is not a duplicate of issue #56833. The provided code snippet triggers an error only when padding="valid" is used, not when padding="same". Additionally, in line with the details outlined in #56833, I have already configured the data_format parameter. @cheyennee,
I was able to replicate the reported behaviour with TF2.14 and tf-nightly as well. Attached [gist](https://colab.research.google.com/gist/tilakrayal/c6b3af27249daeee7c178c4bb05628e3/untitled1521.ipynb) for reference.
Currently we are investigating the issue & will deep dive into the issue and provide the root-cause for the same. Thank you!
### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

No

### Source

source

### TensorFlow version

tf 2.14.0

### Custom code

Yes

### OS platform and distribution

windows colab

### Mobile device

_No response_

### Python version

_No response_

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

Forward AD threw error in tf.autodiff.ForwardAccumulator for tf.keras.layers.AveragePooling3D, but backward AD succeeded with same input.

### Standalone code to reproduce the issue

```shell
import tensorflow as tf

pool_size = [2,1,1]
strides = [2,2,2]
padding = "valid"
data_format = "channels_last"
input = tf.constant(0.895205,shape=[1,1,1,1,1], dtype=tf.float32)
layer = tf.keras.layers.AveragePooling3D(pool_size=pool_size, strides=strides, padding=padding, data_format=data_format, )

with tf.GradientTape(persistent=True, ) as g:
    g.watch(input)
    res_backward = layer(input)
grad_backward = g.jacobian(res_backward,res_backward)
print("res_backward:",res_backward)
print("grad_backward:",grad_backward)

tangents = tf.constant(1.,dtype=tf.float32,shape=[1,1,1,1,1])
with tf.autodiff.ForwardAccumulator(input, tangents) as acc:
    res_forward = layer(input)
    grad_jvp = acc.jvp(res_forward)
    print("res_forward:", res_forward)
    print("grad_forward", grad_jvp)
```


### Relevant log output

```shell
res_backward: tf.Tensor([], shape=(1, 0, 1, 1, 1), dtype=float32)
grad_backward: tf.Tensor([], shape=(1, 0, 1, 1, 1, 1, 0, 1, 1, 1), dtype=float32)
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-3-c1962478ba56> in <cell line: 18>()
     17 tangents = tf.constant(1.,dtype=tf.float32,shape=[1,1,1,1,1])
     18 with tf.autodiff.ForwardAccumulator(input, tangents) as acc:
---> 19     res_forward = layer(input)
     20     grad_jvp = acc.jvp(res_forward)
     21     print("res_forward:", res_forward)

1 frames
/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/ops.py in _create_c_op(graph, node_def, inputs, control_inputs, op_def, extract_traceback)
   1019   except errors.InvalidArgumentError as e:
   1020     # Convert to ValueError for backwards compatibility.
-> 1021     raise ValueError(e.message)
   1022 
   1023   # Record the current Python stack trace as the creating stacktrace of this

ValueError: Exception encountered when calling layer 'average_pooling3d_2' (type AveragePooling3D).

Negative dimension size caused by subtracting 2 from 1 for '{{node gradient_tape/gradient_tape/AvgPool3D}} = AvgPool3D[T=DT_FLOAT, data_format="NDHWC", ksize=[1, 2, 1, 1, 1], padding="VALID", strides=[1, 2, 2, 2, 1]](tangents)' with input shapes: [1,1,1,1,1].

Call arguments received by layer 'average_pooling3d_2' (type AveragePooling3D):
  • inputs=tf.Tensor(shape=(1, 1, 1, 1, 1), dtype=float32)
```
