def ConvertAndCompare(self, func_jax: Callable, *args, with_function: bool=False, atol=None, rtol=None) -> Tuple[(Any, Any)]:
    'Compares jax_func(*args) with convert(jax_func)(*args).'
    func_tf = jax2tf.convert(func_jax)
    if with_function:
        func_tf = tf.function(func_tf)
    res_jax = func_jax(*args)
    res_tf = func_tf(*args)
    self.assertAllClose(res_jax, res_tf, atol=atol, rtol=rtol)
    return (res_jax, res_tf)
Compares jax_func(*args) with convert(jax_func)(*args).
def ConvertAndCompare(self, func_jax: Callable, *args, with_function: bool=False, atol=None, rtol=None) -> Tuple[(Any, Any)]:
    
    func_tf = jax2tf.convert(func_jax)
    if with_function:
        func_tf = tf.function(func_tf)
    res_jax = func_jax(*args)
    res_tf = func_tf(*args)
    self.assertAllClose(res_jax, res_tf, atol=atol, rtol=rtol)
    return (res_jax, res_tf)