DOC: clarify behavior of lax.cond & lax.select

[jakevdp]

In response to questions at #13586; also #7934

Related to #8409

[skirsten]

I also just found #7934 which will probably be resolved by this.

[jakevdp]

Thanks for the pointer!

[jakevdp]

@froystig – Updated based on offline discussion - please take a look!

[bheijden]

I've been noticing super linear compilation times in my application when using jax.lax.switch with an increasing number of branches (similar to this issue). From what I read here, jax.lax.cond is implemented as a special case of switch.

Could these lines explain the increased compilation time? You state that XLA may nevertheless choose to use a select instead of branching if it is deemed advantageous, but could this process of inference lead to super linear compilation times?

In my application, I chose to use a switch instead of a select because I feared long compilation times, but your comments seem to suggest that under the hood XLA may be ignoring this. If this is indeed the case, it would be great if there was an option to force jax/xla to forgo on this optimization to lower the compilation times.