Do jax get rid of unused branches

[jakubMitura14]

Hello I have a model with multiple hyperparameters - some parts of the code should be executed and some not ; the information wheather those parts will be executed or not is known at compile time. For example:

def a(x):
   return x+1
def b(x):
   return x+1

class Model(nn.Module):
   cfg:ml_collections.config_dict.FrozenConfigDict
   def __call__(self,x):
      if(cfg.is_a):
           return a(x)
      return b(x)

As in the toy example above information is in the immutable frozen dict and known in compile time.
Now in real model both function a and b are quite costly. What is additional problem model is vmapped - so using jax lax cond will be automatically change to jax.lax.select hence both branches would be executed.

So in this scenario what is best use if, lax.select; something else, is there a way to point jit that this part of the code for this configuration is a dead end?

[jakevdp]

In general, the semantics of lax.cond allow the compiler to choose not to execute the unneeded branch. lax.select generally leads to evaluation of both branches. Note that if lax.cond is vmapped, it will be lowered to lax.select, and in general both branches will execute.

From your question it's not clear whether any of this is relevant, because you have a static, trace-time condition (using a Python if statement). In that case, only one branch will be traced and therefore only one branch will run.

[jakubMitura14]

hello thanks @jakevdp for answer! yes the model is vmapped - so as far as I understand for static conditions when the model is vmapped old python if statement is the optimal one ?

[jakevdp]

Yes, if you have static conditions you should use a Python if statement, and then the compiler won't even see the other branch.