Allow freezing only subset of data and dims

[ricardoV94]

Description

This improves the optimization introduced on #7047, to allow freezing a subset of dims and data It also fixes a bug where the dim_lengths were not returned

Checklist

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

Type of change

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

---

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

[ricardoV94]

One day mypy will be useful, today is not that day

[ricardoV94]

We should have a test that the original model can still be changed. Otherwise looks fine to me!

Added that check to the first test