[Pallas] Use jnp.where() for branches with tensor-derived conditions but no in-place tensor updates

[AmesingFlank]

Stacked PRs:

• ->[Pallas] Use jnp.where() for branches with tensor-derived conditions but no in-place tensor updates #1936
• Refactor if-else branches, unblock test_if_new_variable_in_static_range for Pallas TPU #1935
• [Pallas] Use fori_loop if loop bounds includes non-constexpr symbolic values #1927

---

[Pallas] Use jnp.where() for branches with tensor-derived conditions but no in-place tensor updates

[v0i0]

huh why does this matter? tensor writes should be vmem mutations, which should be supported?

[AmesingFlank]

The issue is that the condition in 'lax.cond' must be a scalar value. If it is not guaranteed to be a scalar, then we can't use lax.cond. However, if the branches have no side effects (vmem mutations), the we can use jnp.where, which does allow non-scalar conditions.

[v0i0]

when you say scalar, what do you mean? i think even in triton if conditions must be size-1 tensors

[AmesingFlank]

By scalar i specifically mean actual scalars, and not size-1 arrays or 0D arrays. See api doc here https://docs.jax.dev/en/latest/_autosummary/jax.lax.cond.html. This PR is for handling when cond is NOT an actually scalar, in which case lax.cond can't be used, but we can use 'jnp.where' which does allow arrays, provided there's no side effects in the branches

[v0i0]

so by scalar you mean a compile-time bool? surely it works for traced bools, which are 0-dim bool arrays?

[AmesingFlank]

hmm, i think you are right. I was looking at an example (similar to the the one in the test in this PR) where the conditional is an actual size-4 array, but I realized just now that this is because we're now handling tl.grid correctly in Pallas, so the shape of the conditional is wrong.