def check_grads(self, f, args, atol=None, rtol=None, delta=None):
    """Check gradients against finite differences.

    Args:
      f: function to check at ``f(*args)``.
      args: a list or tuple of argument values.
      atol: absolute tolerance for gradient equality.
      rtol: relative tolerance for gradient equality.
      delta: step size used for finite differences.
    """
    if delta is None:
      # Optimal stepsize for central difference is O(epsilon^{1/3}).
      dtype = tf_np.result_type(*args)
      epsilon = onp.finfo(dtype).eps
      delta = epsilon ** (1.0 / 3.0)
    theoretical, numerical = tf.test.compute_gradient(
        to_tf_fn(f), args, delta=delta)
    self.assertAllClose(theoretical, numerical, check_dtypes=False, atol=atol,
                        rtol=rtol)
Check gradients against finite differences.

    Args:
      f: function to check at ``f(*args)``.
      args: a list or tuple of argument values.
      atol: absolute tolerance for gradient equality.
      rtol: relative tolerance for gradient equality.
      delta: step size used for finite differences.
    
trax/tf_numpy/jax_tests/test_util.py