Don't include local_uint_constant_indices rewrite in JAX mode due to XLA bug

[ricardoV94]

Temporarily minimizes #395

Graphs may still fail if users explicitly create uint indexes but that's not very likely, and anyway it's a JAX bug.

[jessegrabowski]

Looks good to me, it's certainty the simplest stop-gap. Do you think it would be worth it to instead temporarily disallow uint8, and make uint16 the smallest allowed? It seems like it would just need to add a check here:

if dtype == np.dtype('uint8'):
    dtype = np.dtype('uint16')

[ricardoV94]

AFAICT this issue happens for every uint:

import jax
import jax.numpy as jnp
import numpy as np

def subtensor(x):
    return x[np.array(65534, dtype="uint16")]

subtensor(jnp.arange(100_000))  # Array(99999, dtype=int32)

I don't think it's worth the trouble. We just wait for JAX to fix it.

We would be penalizing other backends as well

[codecov-commenter]

Codecov Report

Merging #400 (0312140) into main (14d2454) will not change coverage.
The diff coverage is n/a.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #400   +/-   ##
=======================================
  Coverage   80.47%   80.47%           
=======================================
  Files         156      156           
  Lines       45556    45556           
  Branches    11161    11161           
=======================================
  Hits        36660    36660           
  Misses       6694     6694           
  Partials     2202     2202           

Files Changed	Coverage Δ
pytensor/compile/mode.py	84.40% <ø> (ø)