Are subfunctions considered pure?

[inversecrime]

I was wondering whether or not the following code is expected to work.
Basically, it comes down to whether or not jax recognizes if a subfunction uses some value from the superfunction.

Currently, it is working. Is this officially supported, maybe even documented? I couldn't find anything.

def f(x):  # should return |x|
    def path_1():
        return x

    def path_2():
        return -x

    return lax.cond(x > 0, path_1, path_2)


print(jax.make_jaxpr(f)(0))

x = jnp.linspace(-5, 5, 1000)
y = jax.jit(jax.vmap(f))(x)

plt.plot(x, y)
plt.show()

[froystig]

Yes – jax.lax.cond supports closure in the conditional's branch functions.

[inversecrime]

Thanks, good to know, this is very convenient.