Interesting while-loop use-case : jacrev Ok, JIT not => scan

[jecampagne]

Hi,
here I think an interesting simple exo to work to a while-loop transformation (idea adapted from this discussion.

Let
$$
X = (x_0,x_2,\dots,x_{n-1})\in [0,1]^n, \quad \sum_{i=0}^{n-1} x_i =  1
$$
Goal : set the highest value of $(x_i)_i$ to $0$ while $\sum_{i=0}^{n-1} x_i> 0.5$

One can for instance code like that using a while-loop:

def func(x):
  x  = x/jnp.sum(x)

  while jnp.sum(x)>0.5:
    idx = jax.lax.argmax(x,0,index_dtype=jnp.int16)
    x = x.at[idx].set(0.)
  return x

and

x = jax.random.uniform(jax.random.PRNGKey(10), shape=(10,))
print(x)
y = func(x)
y, jnp.sum(y)

gives

0.63337177 0.50169736 0.65120645 0.12358142 0.96813596 0.83877037
 0.91412169 0.60858754 0.88261905 0.71854831]

(Array([0.09258955, 0.07334071, 0.09519672, 0.01806577, 0.        ,
        0.        , 0.        , 0.08896646, 0.        , 0.1050411 ],      dtype=float64),
 Array(0.4732003, dtype=float64))

Interestingly the jacobian FORWARD & BACKWARD can be computed as while is not jax.lax.while_loop

jacfwd_f= jax.jacfwd(func)(x)
np.alltrue(jax.jacrev(func)(x) == jacfwd_f)  # gives True  

BUT, the JIT crashes due to

ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(bool[])>with<DynamicJaxprTrace(level=1/0)>
The problem arose with the `bool` function. 
----->   while jnp.sum(x)>0.5:

So, I have modified the code as followed, which shows that one should think different using JAX if one wants the full power

def jfunc_bis(x):
  n = x.size
  x  = x/jnp.sum(x)

  def noop(carry):
    return carry

  @jit
  def step_once(carry):
    (step, sum, x) = carry
   
    # this part was the core of the while loop
    idx = jax.lax.argmax(x,0,index_dtype=jnp.int16)
    x = x.at[idx].set(0.)

    carry = (step+1, jnp.sum(x), x)
    return carry

  @jit
  def body_fn(carry, dummy):
      (step, sum, x) = carry
      return jax.lax.cond(sum < 0.5, noop, step_once, carry), None   # the while condition is here


  carry = (0, 1., x)
  carry, _ = jax.lax.scan(body_fn, carry, None, length=n)  
  return carry[2]

Then one finds the same results of the non-Jitted version

x = jax.random.uniform(jax.random.PRNGKey(10), shape=(10,))
print(x)
y = jfunc_bis(x)
y

np.alltrue(jax.jacfwd(jfunc_bis)(x) == jacfwd_f)   # here jacfwd_f is the result of non-JIT version
np.alltrue(jax.jacrev(jfunc_bis)(x) == jacfwd_f)

Takeaway:

• While transformation is needed :
  • here due to JIT which is depend value (the * in the doc table)
• the scan version fully all the requirements but needs a bit of work

Now, may be my jfunc_bis is not the best way, so if someone has a better solution (not for the given exemple but for a possible schema of generalisation in other use-case) I will be glad too.

PS: In fact at the start of the exo I was trying to set a jax.lax.while_loop code to make the jacrev crash...

@jit
def new_func(x):
  x  = x/jnp.sum(x)

  def cond(val):
    sum, _ = val
    # if sum>0.5: return True crash due to ConcretizationTypeError even w/o JIT
    return jnp.where(sum>0.5, True, False) 
  def body(val):
    sum, x = val
    idx = jax.lax.argmax(x,0,index_dtype=jnp.int16)
    x = x.at[idx].set(0.)
    return jnp.sum(x), x

  val= jax.lax.while_loop(cond, body,(1.,x))

  return val[1]

then the jacrev will then crash as expected.

[jakevdp]

Thanks for sharing! This is similar in spirit to the "early exit scan" or "bounded while loop" proposed here: #13062

[jecampagne]

ho! @jakevdp do you mean that the jax.lax.scan will have new features in future release (I'm on the 0.4.4)

[jakevdp]

These are just proposed changes at the moment, I don't think there's any definitive timeline for landing them. Just wanted to point out that other people are thinking about the same issues you're thinking about.