fix derived scan logprob when observed provides more broadcastable information

[eclipse1605]

Description

fixed a failure in derived scan logprob construction when the observed/value tensor provides more static broadcastability information than the generative scan graph (e.g. observed has a size-1 axis like (date, 1) while the scan state was inferred as non broadcastable on that axis).

in this, model.logp() could fail during the measurable scan rewrite with a scan outputs_info broadcast pattern mismatch (scan output inferred as matrix like vs. outputs_info expecting vector-like).

• the scan logprob rewrite propagate/align broadcastability metadata introduced by observed/value variables so the rewritten scan is internally consistent.
• this happens without turning scan inner placeholders into non nominal Apply nodes (so scan reconstruction remains valid).

note

I think the same idea can be generalized by treating static broadcastability metadata as part of the measurable scan rewrite contract:

• identify the “outer” value/observed variables that participate in the measurable rewrite and extract their broadcastable axes (ignoring the time axis for sequences).
• when creating inner scan placeholders / outputs_info proxies for the logprob rewritten scan, ensure their TensorType.shape reflects any size-1/broadcastable axes implied by the outer variables.
• apply the same normalization to tapped init buffers (pt.join / outputs_info) so that init and scan outputs agree on broadcastability, without inserting broadcast Apply nodes into the inner graph (placeholders must remain nominal vars).

Related Issue

• Closes #7892

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

[ricardoV94]

I'm inclined to fix this on the Scan side in PyTensor. Basically if inp.type.filter_variable(out.type) passes we should be good to go. This would mean that shape=(None,) -> (5,) or (None,)->(1,) is fair game, just not (5,)->(2,) or something that changes number of dims.

Furthemore shape=(5)->(None,) would end up with a specify_shape in the output and therefore shape=(5,)->(5,). This is the same rule we use for rewrites.

Otherwise it's too tricky to work with scans programatically, and we end up with failures/awkward work-arounds as demonstrated in this PR

[eclipse1605]

sorry for the slow reply, am i getting this right, implementation wise? we’d update PyTensor so that:

• in check_broadcast, we only raise if the inner expects broadcastable axes that the outer doesn’t have, but we allow the outer to be more broadcastable/specific than the inner.
• in Scan.validate_inner_graph, we consider input/output types compatible if the input type is a supertype (same dtype/ndim, output can be more specific), rather than requiring exact broadcastable equality.

[eclipse1605]

@ricardoV94 does this sound right?

[ricardoV94]

The inner should have as much information as the outer, the outre can have less. But yes. Let's open a PR and see how it goes

[eclipse1605]

so in pytensor/pytensor/scan/op.py

we change if b1 != b2 to if b1 and not b2

and in validate_inner_graph instead of

if (
    type_input.dtype != type_output.dtype
    or type_input.broadcastable != type_output.broadcastable
):

we replace the broadcastable eq check with a compatibility check using TensorType.is_super()

if type_input.dtype != type_output.dtype:
    # dtype error
elif isinstance(type_input, TensorType) and isinstance(type_output, TensorType):
    if not type_input.is_super(type_output):
        # compatibility error using is_super
else:
    if type_input != type_output:
        # fallback error

[ricardoV94]

I think you can use is_super always, but let's try

[eclipse1605]

@ricardoV94 i have openned a PR (pymc-devs/pytensor#1861)

[eclipse1605]

@ricardoV94 should i close this PR?

[ricardoV94]

No, but it should be changed to just have a regression test for the original issue to make sure it's working once we fix it in PyTensor

[eclipse1605]

yeah makes sense