Nested -> flat cond()

[MRiabov]

Hello Jax,
It's something that I didn't consider myself, but an LLM has produced such a hallucination:

return cond(
            predicate_1 == 10,
            function_1,
            predicate_1 == 20 & predicate_2 == 20,
            function_2,
            predicate == 30,
            function_3,
            predicate == 40,
            function_4,
            function_5,
        )

A flat cond that accomodates multiple statements. This is an LLM halucination, but it actually makes for a readable code.
I'm aware of switch, but it oftentimes wouldn't work

Could we, pretty please, add nested functionality into cond()? Because the canonical way is just too lengthy:

# snippet from actual code

def select_type(type):
    return lax.cond(
        type == 0,
        lambda _: body_body,  # case when type == 0
        lambda _: lax.cond(
            type == 1,
            lambda _: shaft_shaft,  # case when type == 1
            lambda _: lax.cond(
                type == 2,
                lambda _: shaft_body,  # case when type == 2
                lambda _: lax.cond(
                    type == 3,
                    lambda _: shaft_shaft_shaft,  # case when type == 3
                    lambda _: shaft_shaft_body  # case when type doesn't match 0, 1, 2, or 3
                )
            )
        )
    )

[jakevdp]

You mentioned lax.switch: I think this is essentially what you asking for, although it switches based on an index rather than a sequence of boolean conditionals. Can you say more about why switch did not work for you?

[jakevdp]

Maybe this is more-or-less what you're looking for?

import jax
import jax.numpy as jnp

def multicond(*args, operands):
  *conditions, default = args[::2]
  functions = (*args[1::2], default)
  conditions = [jnp.asarray(c).astype(bool) for c in conditions]
  assert len(functions) == len(conditions) + 1
  assert all(callable(f) for f in functions)
  assert all(c.shape == () for c in conditions)
  index = jnp.array([*conditions, True]).argmax()
  return jax.lax.switch(index, functions, operands)

print(multicond(False, lambda x: 1, True, lambda x: 2, lambda x: 3, operands=(0,)))
# 2

[MRiabov]

You mentioned lax.switch: I think this is essentially what you asking for, although it switches based on an index rather than a sequence of boolean conditionals. Can you say more about why switch did not work for you?

And what if we have variable conditions every time (which is most often the case)?

Pasting from the real code:

return cond(
        f_tang < (mu * gamma_n), # condition one
        inside_upper_cone,
        lambda: cond(
            (f_tang < -(1 / mu) * gamma_n) | (jnp.abs(gamma_n) < 1e-15), # completely unrelated condition two. Switch can't evaluate it, or 
            inside_lower_cone,
            project_to_cone
        )
    )

Why write the extra lambda:cond() here when the code can handle it? And if there 3-4 nested conds it becomes a mess very fast.

OK, there are weird workarounds with switch, namely...

condtions=[lambda x: x*2=4, lambda x: log(x)=1, lambda x: x * sum(1,2,3)=24]
argpos=jnp.where(... (find where hit true) )
result = lax.switch(argpos, [func1,func2,func3], args]

But this is punk, and writing this in code defies the purpose of being consise. Oh, wait, that's basically what you suggested. Well, this should be in some kind of util.

Maybe this is more-or-less what you're looking for?

Well, yes, but why do we not have this in the standard lib? I'm sure many people use nested conds, in RL environments for example.

[jakevdp]

why do we not have this in the standard lib? I'm sure many people use nested conds

Well, this is the first time I've heard that functionality requested, and it can be implemented in a few lines if you need it. I'm not sure it warrants creation of a new API at this point.