pymc-devs / pymc-devs/pytensor

Implement IfElse with inner functions so it works in all linkers

Open
#501 6 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

backend compatibility jax numba
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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.