Repository: tensorflow
Issue: Inexact numeric jacobian causes test failures
Labels: stat:awaiting tensorflower,type:bug,comp:ops,TF 2.9

Using TensorFlow 2.4 with GPUs

**Describe the current behavior**

//tensorflow/python/keras/integration_test:gradients_test fails when run on a node with GPUs. I traced the problem to inexact computation in `_compute_numeric_jacobian`

Output is:
```
FAIL: testLSTMBatchJacobian (__main__.GradientsTest)
GradientsTest.testLSTMBatchJacobian
...
AssertionError: 
Not equal to tolerance rtol=0.01, atol=1e-06
Mismatched value: a is different from b. 
not close where = (array([0]), array([0]), array([2]))
not close lhs = [0.00074506]
not close rhs = [0.00076706]
not close dif = [2.20043e-05]
not close tol = [8.670623e-06]
dtype = float32, shape = (1, 1, 6)
Mismatch: 16.7%
Max absolute difference: 2.20043e-05
Max relative difference: 0.02868646
 x: array([[[-0.013396,  0.007078,  0.000745, -0.02031 ,  0.010461,
         -0.004366]]], dtype=float32)
 y: array([[[-0.0134  ,  0.007075,  0.000767, -0.020303,  0.010446,
         -0.004386]]], dtype=float32)
```

The numeric result is `array([-0.013396,  0.007078,  0.000745, -0.02031 ,  0.010461, -0.004366], dtype=float32)`
while the eager_result, function_result and backprop_result all are `y: array([-0.0134  ,  0.007075,  0.000767, -0.020303,  0.010446, -0.004386], dtype=float32)`


**Standalone code to reproduce the issue**
Reduced extracted test code:

```
import numpy as np
import tensorflow as tf

class GradientsTest(tf.test.TestCase):
  def testLSTMBatchJacobian(self):
    class HasLSTM(tf.keras.Model):

      def __init__(self):
        super(HasLSTM, self).__init__()
        self.lstm = tf.keras.layers.LSTM(units=5)
        self.dense = tf.keras.layers.Dense(1, activation=tf.nn.sigmoid)

      def call(self, x):
        return self.dense(self.lstm(x))

    m = HasLSTM()

    def jacobian(x):
      with tf.GradientTape() as tape:
        tape.watch(x)
        y = m(x)  # pylint: disable=not-callable
      return tape.batch_jacobian(y, x)

    inp = tf.nn.l2_normalize(tf.ones([1, 2, 3]), axis=[1, 2])
    eager_result = jacobian(inp)
    #function_result = tf.function(jacobian)(inp)
    #self.assertAllClose(eager_result, function_result)
    backprop_result, numeric_result = tf.test.compute_gradient(
        m, [inp], delta=1e-3)

    self.assertAllClose(numeric_result, backprop_result, rtol=1e-2)
    self.assertAllClose(tf.reshape(numeric_result, [-1]),
                        tf.reshape(eager_result, [-1]), rtol=1e-2)

if __name__ == "__main__":
  tf.test.main()
```


**Other info / logs**
Increasing the delta to `1e-2` makes the test pass.

Side note: Uncommenting the `tf.function` line yields a couple wrong-looking errors/warnings:
```
WARNING:tensorflow:Using a while_loop for converting CudnnRNNBackprop
pfor.py:1052] Using a while_loop for converting CudnnRNNBackprop
I tensorflow/compiler/mlir/mlir_graph_optimization_pass.cc:116] None of the MLIR optimization passes are enabled (registered 2)
E tensorflow/core/grappler/optimizers/meta_optimizer.cc:592] function_optimizer failed: Invalid argument: Input 0 of node loop_body/PartitionedCall/pfor/PartitionedCall/gradients/zeros_like/pfor/ZerosLike was passed float from has_lstm/lstm/PartitionedCall:6 incompatible with expected variant.
E tensorflow/core/grappler/optimizers/meta_optimizer.cc:592] function_optimizer failed: Invalid argument: Input 0 of node loop_body/PartitionedCall/pfor/PartitionedCall/gradients/zeros_like/pfor/ZerosLike was passed float from has_lstm/lstm/PartitionedCall:6 incompatible with expected variant.
 W tensorflow/core/common_runtime/process_function_library_runtime.cc:805] Ignoring multi-device function optimization failure: Invalid argument: Input 0 of node loop_body/PartitionedCall/pfor/PartitionedCall/gradients/zeros_like/pfor/ZerosLike was passed float from has_lstm/lstm/PartitionedCall:6 incompatible with expected variant.
```