is JAX suitable for if-else heavy program?

[Emerald01]

Hi,

I am new to JAX. I find JAX interesting is because I feel JAX numpy array might be a potential tool to write if-else workflow in GPU.

With initial trials, it seems at least JAX can reach the same mathematical behavior, however the speed is super slow. I think that might be due to the fact my use case involves a lot of if-else branching? This is unavoidable, but my question is this, in general, when evaluating if-else statement, what JAX is going to do? Is it returning the evaluation from GPU back to CPU like a print statement, if so then I do not think JAX is suitable for my problem.

To make my question more concrete, for example, in the following code, I have a simple step(), it checks if the agent tries to cross the boundary, or if the agent hits the landmark. This is a typical skeleton of a path search code. You can see, it has many (seemingly) unavoidable if-else statements and for loop, I struggled hard to make it @jit compatible, and it can run under VMAP, so it can in parallel bundle multiple runners doing the same thing. But it is super slow (several hundred millisecond for this simple computation) compared with a plain numpy counterpart by a few thousand times.

Since I am new to JAX, I would like to see suggestions on my problem, is JAX suitable and how? If not, any advice to approaching this kind of problem? Thank you.

def jax_action_step(action, state, world_boundary, landmarks):
  
  # check if the agent is going to cross the border
  def check_boundary(state, world_boundary):
    
    state, r = L.cond(state[0] < world_boundary[0], 
                 lambda state: (jax.ops.index_update(state, jax.ops.index[0], world_boundary[0]), -0.5),
                 lambda state: (state, 0.0),
                 state)
    
    state, r = L.cond(state[0] > world_boundary[1], 
                 lambda state: (jax.ops.index_update(state, jax.ops.index[0], world_boundary[1]), -0.5),
                 lambda state: (state, 0.0),
                 state)
    
    state, r = L.cond(state[1] < world_boundary[0], 
                 lambda state: (jax.ops.index_update(state, jax.ops.index[1], world_boundary[0]), -0.5),
                 lambda state: (state, 0.0),
                 state)
    
    state, r = L.cond(state[1] > world_boundary[1], 
                 lambda state: (jax.ops.index_update(state, jax.ops.index[1], world_boundary[1]), -0.5),
                 lambda state: (state, 0.0),
                 state)
    
    return state, r

  state = state + action

  # check if the agent hits any landmark
  with loops.Scope() as s:
    s.i = 0
    s.hit = 0
    s.hit_index = -1
    for _ in s.while_range(lambda: s.i < landmarks.shape[0]):
       ### SOME if-else ...
  
  return state, r, s.hit_index

def jax_step(actions, states, world_boundary, landmarks):
  action_step_spawn = jit(vmap(jax_action_step, (0, 0, None, None), 0))
  action_step_spawn(actions, states, world_boundary, landmarks)
  ..........

actions = N.array([[1, 1], [-1, 1], [0, 1]])
states = N.array([[1, 1], [1, 2], [0, 2]])
world_boundary = N.array([0,5])
landmarks = N.array([[2,2],[3,3]])

 jax_step(actions, states, world_boundary, landmarks)

%timeit -n 5 -r 5 jax_step(actions, states, world_boundary, landmarks)
5 loops, best of 5: 291 ms per loop

[davisyoshida]

It seems like you're calling jax.jit on every step. I think you should call jit and vmap before you begin running your loop, then call the jitted function in the loop.

[Emerald01]

Thank you, but I do not get your main point. In this example, I do not have much loop involved, the heavy part is actually four if-else statement, and I called jit(vmap(jax_action_step, (0, 0, None, None), 0)) to jit and distribute. Do you mean, I should reverse some order right here? Sorry, I did not quite get the point.

[davisyoshida]

My point is that you should not time JIT-ing the function. You should JIT the function once, then call it many times. Executing jax.jit(f) is pretty slow and is the sort of thing you should do at the beginning of your code, not during your main loop.