Skip to content

ENH: general framework for non-overlapping switch logp transforms (monotone branches) #8051

@eclipse1605

Description

@eclipse1605

Description

PR #7995 adds logp support for a restricted non-overlapping switch transform (leaky-ReLU pattern) by gating between branch logps evaluated at the observed value. The next step is to track a broader framework to support switch(cond(x), f_true(x), f_false(x)) where:

  • x is a single continuous measurable RV
  • both branches are measurable (PyMC already knows logp(f_true(x), value) and `logp(f_false(x), value)).
  • The branch maps are monotone on their respective domains and have non-overlapping images, so value deterministically implies which branch was taken

The goal is to extend beyond pure scaling (piecewise linear) to a curated set of monotone transforms where we can safely infer a predicate value_implies_true_branch(value) plus runtime constraints that make the inference valid. The switch logp should remain a thin dispatcher: logp = switch(value_implies_true_branch(value), logp(true_branch, value), logp(false_branch, value))

Examples of target patterns:

  • switch(x > 1, x**3, x) (true branch implies value > 1)
  • switch(x > 0, exp(x) - 1, x) (true branch implies value > 0)
  • switch(x > -1, x, exp(x+1) - 2) (true branch implies value > -1)

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions