def test_xla_cpu_gpu_loop_cond_bug(self):

    def deriv(f):
        return (lambda x, *args: jax.linearize((lambda x: f(x, *args)), x)[1](1.0))

    def _while_loop(cond_fun, body_fun, init_val, max_iter):

        def _iter(val):
            next_val = body_fun(val)
            next_cond = True
            return (next_val, next_cond)

        def _fun(tup, _):
            (val, cond) = tup
            return (jax.lax.cond(cond, _iter, (lambda x: (x, False)), val), _)
        init = (init_val, cond_fun(init_val))
        return jax.lax.scan(_fun, init, None, length=max_iter)[0][0]

    def my_pow(x, y):

        def body_fun(val):
            return (val * x)

        def cond_fun(val):
            return True
        return _while_loop(cond_fun, body_fun, 1.0, y)
    self.assertAllClose(deriv(my_pow)(3.0, 1), 1.0, check_dtypes=False)
def test_xla_cpu_gpu_loop_cond_bug(self):

    def deriv(f):
        return (lambda x, *args: jax.linearize((lambda x: f(x, *args)), x)[1](1.0))

    def _while_loop(cond_fun, body_fun, init_val, max_iter):

        def _iter(val):
            next_val = body_fun(val)
            next_cond = True
            return (next_val, next_cond)

        def _fun(tup, _):
            (val, cond) = tup
            return (jax.lax.cond(cond, _iter, (lambda x: (x, False)), val), _)
        init = (init_val, cond_fun(init_val))
        return jax.lax.scan(_fun, init, None, length=max_iter)[0][0]

    def my_pow(x, y):

        def body_fun(val):
            return (val * x)

        def cond_fun(val):
            return True
        return _while_loop(cond_fun, body_fun, 1.0, y)
    self.assertAllClose(deriv(my_pow)(3.0, 1), 1.0, check_dtypes=False)