feat(jax): Allow boolean masks with constants

[aseyboldt]

Indexing with boolean arrays is problematic in jax, because it always needs to know the shape of the result. But this is the case if the mask is a constant, so we can allow that.

TODO needs a test...

[codecov-commenter]

Codecov Report

Merging #376 (2218dbe) into main (981be2a) will increase coverage by 0.00%.
The diff coverage is 100.00%.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #376   +/-   ##
=======================================
  Coverage   80.43%   80.44%           
=======================================
  Files         156      156           
  Lines       45447    45472   +25     
  Branches    11125    11137   +12     
=======================================
+ Hits        36555    36579   +24     
  Misses       6687     6687           
- Partials     2205     2206    +1     

Impacted Files	Coverage Δ
pytensor/link/jax/dispatch/subtensor.py	78.68% <100.00%> (-0.98%)	⬇️

... and 1 file with indirect coverage changes

[ricardoV94]

You might need to inline it in the dispatched function, something like what we do for here

pytensor/pytensor/link/jax/dispatch/shape.py

Line 62 in 7218431

if isinstance(shape, Constant):

[ricardoV94]

I think this was addressed by #849