Allow mutable shape in partially imputed variables

[ricardoV94]

Description

Related Issue

• Closes BUG: NotConstantValueError when using coordinates observed data with missing values with pymc==5.14.0 #7304

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--7314.org.readthedocs.build/en/7314/

[jessegrabowski]

Looks straight-forward enough. Seems like the JAX CI failures are unrelated, but I was wondering if it could interact with freeze_rvs_and_dims somehow. Maybe test?

[ricardoV94]

Looks straight-forward enough. Seems like the JAX CI failures are unrelated, but I was wondering if it could interact with freeze_rvs_and_dims somehow. Maybe test?

Seems to be a change in JAX API of clip, not related to this PR

[jessegrabowski]

Yeah I saw. It just got my noggin' joggin' about the other thing.

[ricardoV94]

Yeah I saw. It just got my noggin' joggin' about the other thing.

We have a test for jax samplers with imputed variables. The question here I guess is freeze_rv_and_dims + direct JAX backend on model with imputed variables. I'll open an issue to add such a test