Derive logprob of < and > operations

[shreyas3156]

This PR implements the logprob inference for binary comparison Ops > and <.

It creates a MeasurableComparison variable and evaluates the logprob based on the truth value of the condition.

Addresses #6633.

Checklist

• Explain important implementation details 👆
• Make sure that the pre-commit linting/style checks pass.
• Link relevant issues (preferably in nice commit messages)
• Are the changes covered by tests and docstrings?
• Fill out the short summary sections 👇

New features

• Logprob inference for GT and LT Ops.

---

📚 Documentation preview 📚: https://pymc--6662.org.readthedocs.build/en/6662/

[codecov]

Codecov Report

Merging #6662 (19ab8ad) into main (5d68bf3) will increase coverage by 0.01%.
The diff coverage is 98.24%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #6662      +/-   ##
==========================================
+ Coverage   91.97%   91.99%   +0.01%     
==========================================
  Files          94       95       +1     
  Lines       15942    16000      +58     
==========================================
+ Hits        14663    14719      +56     
- Misses       1279     1281       +2     

Impacted Files	Coverage Δ
pymc/logprob/binary.py	98.00% <98.00%> (ø)
pymc/logprob/__init__.py	100.00% <100.00%> (ø)
pymc/logprob/transforms.py	95.67% <100.00%> (-0.21%)	⬇️
pymc/logprob/utils.py	100.00% <100.00%> (ø)

... and 3 files with indirect coverage changes

[ricardoV94]

This is awesome!

I left some comments below, but I think it's nearly there.

[shreyas3156]

@ricardoV94 I have made the suggested changes and I just wanted to confirm the approach before adding a test for the separate function logprob/utils.py. (As mentioned in a comment above)

[ricardoV94]

Looks good! Wanna tackle >= and <= next?

[shreyas3156]

Looks good! Wanna tackle >= and <= next?

Definitely! Although I was wondering if >= and <= can be canonicalized to LT and GT respectively?

[ricardoV94]

Looks good! Wanna tackle >= and <= next?

Definitely! Although I was wondering if >= and <= can be canonicalized to LT and GT respectively?

Yes if you allow either direction in const<dist and const>dist. In this PR only dist<const and dist>const works no?

[shreyas3156]

Yes if you allow either direction in const<dist and const>dist. In this PR only dist<const and dist>const works no?

We can make this PR process expressions like const<dist by making the node rewriter check which input holds the base_rv. We can then assign the outputs to the rewritten node in a fixed order for all the comparison operations. (e.g. (base_rv, const))

But this may require adding a transform for <= and >= to GT and LT.

I think it may be better to simply add separate derivations of LE and GE like in this PR.

[ricardoV94]

Yes if you allow either direction in const<dist and const>dist. In this PR only dist<const and dist>const works no?

We can make this PR process expressions like const<dist by making the node rewriter check which input holds the base_rv. We can then assign the outputs to the rewritten node in a fixed order for all the comparison operations. (e.g. (base_rv, const))

But this may require adding a transform for <= and >= to GT and LT.

I think it may be better to simply add separate derivations of LE and GE like in this PR.

You can see that for other MeasurableElemwise like Add we store a measurable_input_idx, so that we can have the right logp later on, regardless of the order in which the dist and consts came in. So maybe something like that would be simpler?

pymc/pymc/logprob/transforms.py

Line 348 in 1ed4475

measurable_input_idx: int

[shreyas3156]

You can see that for other MeasurableElemwise like Add we store a measurable_input_idx, so that we can have the right logp later on, regardless of the order in which the dist and consts came in. So maybe something like that would be simpler?

Sure, I'll look into this approach.

Meanwhile, I had a question about the expression used in this PR to calculate logprob for discrete distributions. I think I should correct it to the following (please correct me if I've misunderstood)

1. For <: CDF - PMF (x = const) when True; 1 - CDF + PMF (x = const) when False
2. For >: 1-CDF when True; CDF when False

[ricardoV94]

You can see that for other MeasurableElemwise like Add we store a measurable_input_idx, so that we can have the right logp later on, regardless of the order in which the dist and consts came in. So maybe something like that would be simpler?

Sure, I'll look into this approach.

Meanwhile, I had a question about the expression used in this PR to calculate logprob for discrete distributions. I think I should correct it to the following (please correct me if I've misunderstood)

1. For `<`: `CDF - PMF (x = const) when True; 1 - CDF + PMF (x = const) when False`

2. For `>`: `1-CDF when True; CDF when False`

For discrete variables the CDF gives us the probability <=, so yes. You can do
logp(dist < const, True) = logcdf(dist, const - 1).

[ricardoV94]

Flagging it so we don't merge with the wrong logp by accident

[ricardoV94]

Looks great

[shreyas3156]

You can see that for other MeasurableElemwise like Add we store a measurable_input_idx, so that we can have the right logp later on, regardless of the order in which the dist and consts came in. So maybe something like that would be simpler?

I went through transforms.py and could not understand why using the index is required. We already extract the measurable_input from the list of inputs first.

pymc/pymc/logprob/transforms.py

Lines 556 to 562 in 61be336

	measurable_inputs = [
	inp
	for idx, inp in enumerate(node.inputs)
	if inp.owner
	and isinstance(inp.owner.op, MeasurableVariable)
	and inp not in rv_map_feature.rv_values
	]

We then make a MeasurableTransform node from transform_inputs which always has measurable_input as its first element, which is also why the measurable_input_idx always seems to have the value 0.

pymc/pymc/logprob/transforms.py

Lines 585 to 586 in 61be336

	measurable_input_idx = 0
	transform_inputs: Tuple[TensorVariable, ...] = (measurable_input,)

[ricardoV94]

You're right, so maybe we don't need it at all? Just make it so the first input is always the measurable one?