Implement proper JAX/Numba dispatch for IfElse

[ricardoV94]

Description

IfElse is only lazy in the default backend because the function virtual machine handles it (via the "lazy" attribute"). In Numba/JAX it currently does nothing, because it receives all outputs pre-computed.

During compilation we could specialize IfElse into a LazyIfElse Op that contains two inner Graphs, one corresponding to each branch. These graphs should contain all variables that lead to the inputs of IfElse and are not used by any other output variable other than through the outputs of IfElse. This depends on which function is being compiled and can't be known ahead of time.

The current implementation of jax_funcify_IfElse:

pytensor/pytensor/link/jax/dispatch/basic.py

Lines 60 to 70 in 4235ccc

	@jax_funcify.register(IfElse)
	def jax_funcify_IfElse(op, **kwargs):
	n_outs = op.n_outs
	def ifelse(cond, *args, n_outs=n_outs):
	res = jax.lax.cond(
	cond, lambda _: args[:n_outs], lambda _: args[n_outs:], operand=None
	)
	return res if n_outs > 1 else res[0]
	return ifelse

Would instead look something like (pseudo-code):

@jax_funcify.register(LazyIfElse)
def jax_funcify_LazyIfElse(op, **kwargs):
    true_fn = jax_funcify(op.true_fgraph)
    false_fn = jax_funcify(op.false_fgraph)

    def ifelse(cond, *args):
        res = jax.lax.cond(cond, true_fn, false_fn, *args)
        return res if n_outs > 1 else res[0]

    return ifelse

This could even provide a nicer dprint, by showing the two inner graphs. Right now it's not always obvious what operations are lazily computed or not.

[Dhruvanshu-Joshi]

Hi
Here's my understanding of what I think should be done:
Create a LazyIfElse op inheriting from the IfElse Op wherein we have seperate subgraphs for true and false conditional output essentially in the self.true_fgraph and self.false_fgraph.
Next we'll have to create a jax_funcify_LazyIfElse as you suggested.

Here's what I think should be the next pseudo-code:

def specialize_ifelse_to_lazy_ifelse(fgraph):
    for node in fgraph.apply_nodes:
        if isinstance(node.op, IfElse):
            condition = node.inputs[0]
            true_inputs = node.inputs[1:1 + node.op.n_outs]
            false_inputs = node.inputs[1 + node.op.n_outs:]
            
            outputs = node.outputs

            true_fgraph = create_subgraph(true_inputs, outputs[:node.op.n_outs])
            false_fgraph = create_subgraph(false_inputs, outputs[node.op.n_outs:])

            lazy_ifelse_op = LazyIfElse(true_fgraph, false_fgraph, node.op.n_outs)

            new_node = lazy_ifelse_op.make_node(condition, *node.inputs[1:])
            
            # Replace IfElse node with LazyIfElse node
            fgraph.replace_all([(old, new) for old, new in zip(node.outputs, new_node.outputs)])

Is this correct?
Also I cannot seem to understand the compilation process. How would I make the call for this specialize_ifelse_to_lazy_ifelse during compilation of jax/numba functions only and identify the op?

[ricardoV94]

So the challenge is to find the largest distinct subgraphs for the true and false branches starting from the same set of inputs. This is a fun exercise in graph partitioning.

Here are two examples to illustrate this:

cond = pt.scalar("cond", dtype=bool)

x = pt.scalar("x")
y = pt.exp(x)
sinx = pt.sin(y)
cosx = pt.cos(y)
xcosx = cosx * x
out = ifelse(cond, sinx, xcosx)

In that case we probably want true_fn and false_fn to be a function from (x, y) to sinx/cosx. We start at y because that's common to both branches and has to be computed anyway

def true_func(x, y):
  return pt.sin(y)

def false_func(x, y):
  return pt.cos(y) * x

So operations that are shared between the two branches should be computed outside and given as inputs. At the same time, we want each function to have the largest number of unique operations inside each branch, to make them truly lazy

One extra complication. If a expression is needed outside of the ifelse graphs as well as inside one of them we don't want to place it inside to avoid recomputing it twice. For example

x = pt.scalar()
out1 = x + pt.exp(x)
out2 = x + pt.log(x)
out = ifelse(x <= 0, out1, out2)
w = out + pt.exp(x)

In this case we don't want to compute exp(x) inside the true branch because we need to compute it anyway outside.

So probably we want the following funs

def true_func(x, y):
  return x + y

def false_func(x, y):
  return x + pt.log(x)

Where y is exp(x)

But actually both branches do x + something so in fact we don't need that operations inside and the graph should first be rewritten as

x = pt.scalar()
out1 = pt.exp(x)
out2 = pt.log(x)
out = x + ifelse(x <= 0, out1, out2) 
w = out + pt.exp(x)

With inner funcs

def true_func(x, y):
  return y

def false_func(x, y):
  return pt.log(x)

However we can treat that issue separately from the main goal of making the op lazy. It doesn't need to be a blocker for an initial solution

Finally, we need an algorithm that naturally handles nested ifelse.

[Dhruvanshu-Joshi]

Hi @ricardoV94
I was playing around with this yesterday for a while. What I did would makes sense more if I open a PR for that. Can you lease review this dprint that I am getting so that we know that I am on the right track:

if{inplace} [id A] 1
 ├─ cond [id B]
 ├─ Composite{...}.0 [id C] 0
 │  └─ x [id D]
 └─ Composite{...}.1 [id C] 0
    └─ ···

Inner graphs:

Composite{...} [id C]
 ← sin [id E] 'o0'
    └─ exp [id F] 't1'
       └─ i0 [id G]
 ← mul [id H] 'o1'
    ├─ cos [id I]
    │  └─ exp [id F] 't1'
    │     └─ ···
    └─ i0 [id G]
    ```

[ricardoV94]

Sounds good, although composite is a distinct op with InnerGraphs restricted to Elemwise operations.

We want LazyIfElse itself two have two inner graphs, like Composite has one, and not restricted to Elemwise operations

You probably can use a lot of the internal code of Composite (or even OpFromGraph which is the most general form of inner function Op)