logprob for non-overlapping switch accepts equivalent spellings

[eclipse1605]

Closes #8049

[codecov]

Codecov Report

❌ Patch coverage is 78.82353% with 18 lines in your changes missing coverage. Please review.
✅ Project coverage is 90.74%. Comparing base (c68c56e) to head (2fe8441).
⚠️ Report is 4 commits behind head on main.

Files with missing lines	Patch %	Lines
pymc/logprob/switch.py	78.82%	18 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #8058      +/-   ##
==========================================
- Coverage   91.42%   90.74%   -0.68%     
==========================================
  Files         117      121       +4     
  Lines       19154    19487     +333     
==========================================
+ Hits        17512    17684     +172     
- Misses       1642     1803     +161     

Files with missing lines	Coverage Δ
pymc/logprob/switch.py	77.85% <78.82%> (-2.83%)	⬇️

... and 12 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[ricardoV94]

I'm hoping we can simplify the code a bit further. Also let's get rid of all these typing.cast. It's in almost every line. Just mark the file as failing mypy if it requires so much work.

[ricardoV94]

Can we simplify this code?

[ricardoV94]

Why not parametrize already with the objects you gonna need?

[ricardoV94]

Instead of having to guess, you could normalize the switch so that it's always switch(cond, x, neg_branch). You can write switch(c, t, f) -> switch(~c, f, t).

Maybe this allows you to simplify more logic elsewhere?

[eclipse1605]

ya makes sense, though we still can’t use cond directly inside the logprob, because cond depends on the latent x

[ricardoV94]

in the logprob you should be able to assume you already have a canonical form, because that's how you made it in the rewrite

[eclipse1605]

actually i think we can do this and it should simplify the logprob. we can canonicalize in the rewrite so downstream logprob can assume one shape.

maybe something like this (pseudocode):

# in measurable rewrite:
# detect either ordering + equivalent conditions
#   switch(cond(x ? 0),  x,   a*x)
#   switch(cond(x ? 0), a*x,  x)
# and also swapped comparisons like 0 < x, 0 >= x, ...

(x, ax, sem) = match_branches_and_condition(node)

# normalize to always: switch(cond_canon, x, ax)
if true_branch_is(ax):
    sem = negate(sem)          # because switch(c, ax, x) == switch(~c, x, ax)
cond_canon = cond_from_semantics(x, sem) 

return measurable_switch_non_overlapping(cond_canon, x, ax)

# then in logprob:
# assume inputs are already canonical: switch(cond_canon, x, a*x)
sem = parse_canonical_cond(cond_canon, x)
gate = value_based_gate(value, sem) 
return switch(gate, logp(x, value), logp(a*x, value)) + check(a > 0)

[ricardoV94]

yeah, let's see if it works

[eclipse1605]

I'm hoping we can simplify the code a bit further. Also let's get rid of all these typing.cast. It's in almost every line. Just mark the file as failing mypy if it requires so much work.

sure, is this a general convention we follow throughout the codebase?

[ricardoV94]

I'm hoping we can simplify the code a bit further. Also let's get rid of all these typing.cast. It's in almost every line. Just mark the file as failing mypy if it requires so much work.

sure, is this a general convention we follow throughout the codebase?

No, but in general I favor not making our code a mess because of mypy

[eclipse1605]

anyways i misunderstood mypy failing locally in the earlier commit, it wasn't required

[eclipse1605]

@ricardoV94 the failing test is flaky right?

[eclipse1605]

@ricardoV94 anything left in this?