Solve for pymc>=5.16.0

[rohanbabbar04]

Closes #482

Description

• Looks like the inputs now return ExpandDims, this is due to the latest release of pytensor==2.23.0

• Made the changes to handle it.

  • [x] Code style is correct (follows pylint and black guidelines)
  • [x] Includes new or updated tests to cover the new feature
  • [x] New features are properly documented (with an example if appropriate)