Do not consider dims without coords volatile if length has not changed

[JasonTam]

Description

Allow coords defined by length only (no values) to be considered constant_coords if the new length matches.

Related Issue

• Closes BUG: Using a length-only coord causes it to be volatile #7376
• Related to #

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):

---

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

[JasonTam]

The main part I feel uneasy about is having to call current_length.eval() since the length stored in model.dim_lengths is a pt shared variable

[ricardoV94]

It can only be a shared variable in which case you can retrieve with get_value(), or a Constant in which case you can retrieve with .data That should be much more efficient than eval unless eval does something clever under the hood

[ricardoV94]

Also we need a test @JasonTam

[JasonTam]

@ricardoV94
I've pulled out the constant coord check logic into its own method in order to test in isolation. As I wasn't sure how to test this against the full sample_posterior_predictive method (or if that's even practical). Also, that method seems to have grown quite large.

[codecov]

Codecov Report

Attention: Patch coverage is 93.33333% with 1 line in your changes missing coverage. Please review.

Project coverage is 92.18%. Comparing base (f719796) to head (b76087d).
Report is 68 commits behind head on main.

Files with missing lines	Patch %	Lines
pymc/sampling/forward.py	93.33%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #7381      +/-   ##
==========================================
- Coverage   92.18%   92.18%   -0.01%     
==========================================
  Files         103      103              
  Lines       17249    17258       +9     
==========================================
+ Hits        15901    15909       +8     
- Misses       1348     1349       +1     

Files with missing lines	Coverage Δ
pymc/sampling/forward.py	95.72% <93.33%> (-0.28%)	⬇️

[ricardoV94]

@JasonTam here is where we test this volatility logic:

pymc/tests/sampling/test_forward.py

Line 113 in cbf1591

class TestCompileForwardSampler:

Can you add something there instead. No need to do actual sampling, just make a scenario where a variable depends on a dim without coords (and then change it's length or not)

[JasonTam]

@ricardoV94 I'm not sure how to write the test without actually sampling. In the similar tests under TestCompileForwardSampler, volatile vars are checked from the output of compile_forward_sampling_function. But that function already takes in constant_coords as a parameter.

The new logic of determining constant_coords can either be tested via the new low level function get_constant_coords (as currently written in the PR), or in the method which calls it, sample_posterior_predictive.

Maybe I'm misunderstanding what you mean.

[ricardoV94]

You're right @JasonTam, I was confused by the fact that the changes in constant data are checked inside the compile_forward_function but the coords are not. I guess this is the case because coords are not part of the graph while SharedVariables are.

[ricardoV94]

Can we add a more integration-like test like?

Spit-balling, perhaps something like:

with pm.Model() as m:
  m.add_coord("trial", length=3)
  x = pm.Normal("x", shape="trial")
  y = pm.Deterministic("y", x.mean())

# pass dims as well  
idata = az.from_dict("x": [[np.pi, np.pi, np.pi]])

with m:
  pp1 = pm.sample_posterior_predictive(idata, var_names=["y"]).posterior_predictive

assert pp1["y"] == np.pi

with m:
  # change coord length 
  pp2 = pm.sample_posterior_predictive(idata, var_names=["y"]).posterior_predictive

assert pp2["y"] != np.pi

[JasonTam]

Does this look ok? d40fffe

requested extra test

[ricardoV94]

I prefer to change the dim length in the pymc model and pass exactly the same idata to sample_posterior_predictive as was done in the first call.

[ricardoV94]

The way to update an existing dim is with Model.set_dim()

[JasonTam]

Sorry, just got to this now. Implemented in 49bbf8c -- hope I understood you correctly

[ricardoV94]

Yup you got it!

[ricardoV94]

Thanks @JasonTam