Inline static size inputs in JAX implementation of RandomVariables

[ricardoV94]

Description

This gets around some limitations in JAX jitting system

This is required for pymc-devs/pymc#7324 to pass

Related Issue

• Closes #
• 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):

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 80.85%. Comparing base (d80c0bf) to head (edb72c2).
Report is 210 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #774   +/-   ##
=======================================
  Coverage   80.85%   80.85%           
=======================================
  Files         162      162           
  Lines       47019    47026    +7     
  Branches    11504    11507    +3     
=======================================
+ Hits        38018    38025    +7     
  Misses       6750     6750           
  Partials     2251     2251           

Files	Coverage Δ
pytensor/link/jax/dispatch/random.py	93.96% <100.00%> (+0.16%)	⬆️

[drbenvincent]

Commenting so I get updates about the progress of this PR :)

[ricardoV94]

I'm converting to draft because I want to see if there's another way to work with constant shapes without having to inline from the dispatch function

[ricardoV94]

Nope this seems to be the easiest way