pymc-devs / pymc-devs/pytensor
Implement IfElse with inner functions so it works in all linkers
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 644
- Forks
- 208
- Avg merge
- 2d 14h
- Merged PRs (30d)
- 16
Description
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:
https://github.com/pymc-devs/pytensor/blob/4235ccc3f4243c5179178a206c15d84c4cda2e79/pytensor/link/jax/dispatch/basic.py#L60-L70
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.
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with the current jax_funcify_IfElse implementation in pytensor/link/jax/dispatch/basic.py and trace how IfElse is compiled for each linker. Determine how compilation can specialize it into a LazyIfElse with true and false inner Graphs containing only branch-specific inputs. Done means lazy evaluation works across linkers and the dprint exposes the two inner graphs.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers
- Issue type
- Feature
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Stale
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100