Guidance on correct way to conditionally apply function

[DanPuzzuoli]

Hi,

I have a list of functions that with signature:

• Input: DeviceArray
• Output: (DeviceArray, DeviceArray) (i.e., they return a pair of device arrays)
  and I want to conditionally apply one based on a list of boolean conditions. (All functions have consistent output type/shapes based on the input array, so no worries there.)

I was hoping to use jnp.piecewise to do this, but I'm running into the error:

AttributeError: 'tuple' object has no attribute 'astype'

so I guess jnp.piecewise must have an array output? I.e. it can't be used for functions that output pytrees?

I alternatively tried modifying the functions so that their output concatenates the two outputs into a single DeviceArray with an extra dimension, but doing this resulted in the error:

ValueError: output shape (2,) does not match core dimensions () on vectorized function with excluded=(1,) and signature=None

where the above error was generated when I tried running the following simple example:

jnp.piecewise(
      x=2., 
      condlist=1. < jnp.array([3.]), 
      funclist=[lambda t: jnp.array([t, t**2]), 
                lambda t: jnp.array([t**2, t**3])]
)

Will I need to more directly use switch to solve this problem? I see that switch can be used on functions whose outputs are Tuples, though I'm slightly confused as to how to extract the index to pass to switch in a compilable way (which I guess is sort of what the content of jnp.piecewise is to begin with?).

Also, as it may be relevant, this question originated when starting to look into finally addressing #5647 - the elimination of some redundant computation in jax.scipy.linalg.expm. The list of functions I want to choose between is the _padex functions, and the conditions will be given by A_L1 < conds.

[jakevdp]

I think jnp.piecewise does require array outputs in order to work correctly, but you should be able to do this directly with lax.cond; for example:

from jax import jit, lax

def f1(x):
  return x, x + 1

def f2(x):
  return x + 1, x + 2

@jit
def switch_f(x):
  return lax.cond(x < 2, f1, f2, x)

print(*switch_f(1))
# 1 2

[DanPuzzuoli]

Thanks!

I should have clarified though that I'm interested in the case where the number of branches may be long, and I am trying to avoid using multiple nested calls to cond.

[jakevdp]

Makes sense. Perhaps then you want lax.switch?

print(*lax.switch(0, [f1, f2], jnp.arange(4)))
# [0 1 2 3] [1 2 3 4]

You can pass this an arbitrary number of functions, along with the index of the function you want to be used.

[soraros]

And the index i can be obtained by

i = jnp.argwhere(jnp.array([c0, c1, c2, ...]), size=1).squeeze()  # where c* are the conditions

[DanPuzzuoli]

This argwhere/squeeze combo is the piece I was missing. Thanks!