nested control flow

[yiminghwang]

Hi, there,

In order to use jit, I have to use jax.lax.cond to implement the condition flow, is there any elegant way to write the nested if?
for example,

original one,

x = 1
y = 0
op = 3

if op == 2:
   if y == 0:
         x *= 1j
   else:
        x *= -1j
elif op == 3:
   if y == 1:
        x *= -1

the version of using lax.cond

x = jnp.complex64(1)
y = 0
op = 3

result = jax.lax.cond(op == 2,
                          lambda a: jax.lax.cond(y == 0,
                                                 lambda b: b * 1j,
                                                 lambda b: b * -1j,
                                                 a),
                          lambda a: jax.lax.cond(op == 3,
                                                 lambda b: jax.lax.cond(y == 0,
                                                                        lambda c: c,
                                                                        lambda c: c * -1,
                                                                        b),
                                                 lambda b: b,
                                                 a),
                          x)

[jakevdp]

The answer depends whether the conditions are static or dynamic. For example, assuming op is static (it sort of looks like a static indicator variable here), I might write it like this:

import jax
from functools import partial

@partial(jax.jit, static_argnames=['op'])
def f(x, y, op=3):
  if op == 2:
    return x * jnp.where(y == 0, 1j, -1j)
  elif op == 3:
    return x * jnp.where(y == 1, -1, 1)
  else:
    raise NotImplementedError(f"{op=} not supported")

[mattjj]

You can also use jax.lax.switch, which makes for a nice dynamic-op analogue to what Jake wrote.

[yiminghwang]

Great, Thank you both!