Derive Categorical from Argmin/Argmax#8072
Conversation
ricardoV94
force-pushed
the
categorical_derived_rvs
branch
2 times, most recently
from
5ff3f21 to
edd4bc7
Compare
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #8072 +/- ##
==========================================
- Coverage 91.73% 91.73% -0.01%
==========================================
Files 128 128
Lines 20697 20807 +110
==========================================
+ Hits 18987 19087 +100
- Misses 1710 1720 +10
🚀 New features to boost your workflow:
|
|
CC @eclipse1605 if you want to review / ask questions |
ricardoV94
force-pushed
the
categorical_derived_rvs
branch
from
edd4bc7 to
36634c0
Compare
|
Check out this pull request on See visual diffs & provide feedback on Jupyter Notebooks. Powered by ReviewNB |
Documentation build overview
19 files changed ·
|
ricardoV94
force-pushed
the
categorical_derived_rvs
branch
from
36634c0 to
8bc8201
Compare
ricardoV94
force-pushed
the
categorical_derived_rvs
branch
2 times, most recently
from
635bd12 to
11eaebb
Compare
ricardoV94
force-pushed
the
categorical_derived_rvs
branch
from
11eaebb to
0717876
Compare
| elif ( | ||
| len(base_node.inputs) == 2 | ||
| and isinstance(base_node.op, Elemwise) | ||
| and isinstance(base_node.op.scalar_op, Mul) | ||
| and ((right_neg := is_minus_1(base_node.inputs[1])) or is_minus_1(base_node.inputs[0])) | ||
| ): | ||
| base_var = base_node.inputs[0 if right_neg else 1] | ||
| base_node = base_var.owner | ||
| if base_node is None: | ||
| return None | ||
|
|
||
| if isinstance(base_node.op, ExponentialRV): | ||
| # argmin(exponential(rate)) -> categorical(rate) |
There was a problem hiding this comment.
wont this break on plain Neg or fused exp * -scale unless lift_loc_scale gets there first? idk...
There was a problem hiding this comment.
Neither is a problem. Neg gets canonicalized to mul by -1, and for the second, bailing is exactly the right thing. The rewrites are recalled until saturation (when a pass changes nothing). So if this comes before the lift_loc_scale, it bails, then lift_loc_scale does its thing, and then the next call is a hit. You can see them happening if you enable config.optimizer_verbose:
import numpy as np
import pytensor
import pytensor.tensor as pt
from pymc.logprob import logp
pytensor.config.optimizer_verbose = True
probs = np.array([0.1, 0.3, 0.6])
x = pt.argmin(pt.random.exponential(scale=1 / probs), axis=-1)
logp(x, x.type())
# rewriting: rewrite local_neg_to_mul replaces Neg.0 of Neg(exponential_rv.out) with Mul.0 of Mul(ExpandDims{axis=0}.0, exponential_rv.out)
# rewriting: rewrite find_measurable_transforms replaces Mul.0 of Mul([-1.], exponential_rv.out) with MeasurableMul.0 of MeasurableMul(exponential_rv.out, [-1.])
# rewriting: rewrite categorical_from_argmax replaces argmax of Argmax{axis=(0,)}(MeasurableMul.0) with categorical_rv.out ...
x = pt.argmax(pt.random.exponential(size=3) * (-1 / probs), axis=-1)
logp(x, x.type())
# rewriting: rewrite lift_loc_scale replaces MeasurableMul.0 of MeasurableMul(exponential_rv.out, [-10. ... .66666667]) with Mul.0 of Mul(exponential_rv.out, ExpandDims{axis=0}.0) # <- lifted rv * -1
# rewriting: rewrite find_measurable_transforms replaces Mul.0 ... with MeasurableMul.0 ...
# rewriting: rewrite categorical_from_argmax replaces argmax of Argmax{axis=(0,)}(MeasurableMul.0) with categorical_rv.out ...There was a problem hiding this comment.
no need to apologize, they were good questions
There was a problem hiding this comment.
thanks :) but for a bit more clarity, the saturation thing makes sense when lift_loc_scale can actually run
for symbolic outer scale i still dont think that rescue happens because lift_loc_scale bails every pass on constant_fold so graph never changes and hits saturation with neither rewrite matching
import numpy as np
import pytensor
import pytensor.tensor as pt
from pymc.logprob import logp
pytensor.config.optimizer_verbose = True
scale = pt.vector("scale")
x = pt.argmin(pt.random.exponential(scale=1, size=3) * scale, axis=-1)
logp(x, x.type())
# rewriting: rewrite local_neg_to_mul replaces Neg.0 of Neg(Mul.0) with Mul.0 of Mul(ExpandDims{axis=0}.0, Mul.0)
# rewriting: rewrite local_mul_canonizer replaces Mul.0 of Mul(ExpandDims{axis=0}.0, Mul.0) with Mul.0 of Mul(ExpandDims{axis=0}.0, exponential_rv{"()->()"}.out, scale)
# rewriting: rewrite find_measurable_transforms replaces Mul.0 of Mul([-1.], exponential_rv{"()->()"}.out, scale) with MeasurableMul.0 of MeasurableMul(exponential_rv{"()->()"}.out, Mul.0)
# >>> NotImplementedError: Logprob method not implemented for Argmax{axis=(0,)}constant outer scale does hit lift + categorical
idk if this makes sense, maybe im missing something
There was a problem hiding this comment.
We reject that because scale could be negative and we can't lift negative scales inside the positive distributions. The inline comment says we could do more work say if scale is abs/exp of a root variable then we'd be safe again
There was a problem hiding this comment.
gotit makes sense thanks :)
2 participants
Couldn't resist the temptation...
Needed some import shuffling so that I'm allowed to import a SymbolicRandomVariable in logprob rewrites