def testIssue1361(self):

    @api.jit
    def jit_run_scan(x):

        def fun(carry, _):
            (x, _) = carry
            return (((2 * x), 0.0), None)
        ((x, _), _) = lax.scan(fun, (x, 0.0), jnp.arange(3))
        return x
    api.grad((lambda x: jit_run_scan(x)))(0.0)
def testIssue1361(self):

    @api.jit
    def jit_run_scan(x):

        def fun(carry, _):
            (x, _) = carry
            return (((2 * x), 0.0), None)
        ((x, _), _) = lax.scan(fun, (x, 0.0), jnp.arange(3))
        return x
    api.grad((lambda x: jit_run_scan(x)))(0.0)