Implement batch evaluation of RV shapes in get_plates

[RohanAich]

Implement batch evaluation of RV shapes in ModelGraph.get_plates

Description

Implements the TODO in ModelGraph.get_plates to evaluate all random variables (RV) shapes at once instead of indiviually.

Related Issue

• Closes #
• Related to TODO in ModelGraph.get_plates

Type of change

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

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 91.44%. Comparing base (056e80c) to head (21b603c).

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #8033   +/-   ##
=======================================
  Coverage   91.44%   91.44%           
=======================================
  Files         116      116           
  Lines       19002    19005    +3     
=======================================
+ Hits        17377    17380    +3     
  Misses       1625     1625           

Files with missing lines	Coverage Δ
pymc/model_graph.py	84.69% <100.00%> (+0.15%)	⬆️

🚀 New features to boost your workflow:

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

[ricardoV94]

Good chance to fix #8024?

We need something better than _cheap_eval_mode, because we need it to do shape inference first to handle flat RV. A possible strategy is to add this helper to pytensorf, and call it before we eval the chapes with the cheap mode:

def get_symbolic_rv_shapes(
    rvs: Sequence[Variable], raise_if_rvs_in_graph: bool = True
) -> tuple[TensorVariable]:
    # TODO: docstrings
    
    rv_shapes = [rv.shape for rv in rvs]
    shape_fg = FunctionGraph(outputs=rv_shapes, features=[ShapeFeature()], clone=True)
    with config.change_flags(optdb__max_use_ratio=10, cxx=""):
        infer_shape_db.default_query.rewrite(shape_fg)
    rv_shapes = shape_fg.outputs

    if raise_if_rvs_in_graph and (overlap := (set(rvs) & set(ancestors(rv_shapes)))):
        raise ValueError(f"rv_shapes still depend the following rvs {overlap}")

    return tuple(rv_shapes)

I would call it with raise_if_rvs_in_graph=False

[RohanAich]

Thank you! This may take me a bit but I will try to get it done within the next few days.