Fix measurable stack and join with interdependent inputs

[ricardoV94]

This fixes measurable stacks and joins that have interdependent variables.

The main loop of factorized_joint_logprob replaces all the measurable variables already seen by their values before calling the logp of a new measurable node, but it does not replace said node by its value variable, as it assumes the returned logprob expression never depends on the original measurable (stochastic) node.

However this can happen in Stack and Join whenever the base measurable RVs depend on one another. The new tests show this clearly. The solution proposed in this PR is to do the additional replacements needed in the respective logp functions.

Major / Breaking Changes

• ...

Bugfixes / New features

• Fix logprob of stack and join when base variables are interdependent

Docs / Maintenance

• ...

[codecov]

Codecov Report

Merging #6342 (ddc6b65) into main (ddc6b65) will not change coverage.
The diff coverage is n/a.

❗ Current head ddc6b65 differs from pull request most recent head 51fbdc9. Consider uploading reports for the commit 51fbdc9 to get more accurate results

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #6342   +/-   ##
=======================================
  Coverage   94.71%   94.71%           
=======================================
  Files         132      132           
  Lines       26695    26695           
=======================================
  Hits        25284    25284           
  Misses       1411     1411           

[michaelosthege]

damn, I forgot to submit my review comments

[michaelosthege]

How do we know that there are no transforms?

[ricardoV94]

The transforms are applied before the RV-specific logp function is called. The value received is already the one in the constrained space.

[ricardoV94]

pymc/pymc/logprob/transforms.py

Lines 591 to 608 in ddc6b65

	@_logprob.register(new_op_type)
	def transformed_logprob(op, values, *inputs, use_jacobian=True, **kwargs):
	"""Compute the log-likelihood graph for a `TransformedRV`.
	We assume that the value variable was back-transformed to be on the natural
	support of the respective random variable.
	"""
	(value,) = values
	logprob = _logprob(rv_op, values, *inputs, **kwargs)
	if use_jacobian:
	assert isinstance(value.owner.op, TransformedVariable)
	original_forward_value = value.owner.inputs[1]
	jacobian = op.transform.log_jac_det(original_forward_value, *inputs)
	logprob += jacobian
	return logprob

[ricardoV94]

Thanks for the review @michaelosthege. Please don't merge this, as I will go through it in the hackathon

[ricardoV94]

CC @tomicapretto