ENH: Add dispatches of TruncatedNormal distribution for forward sampling

[lucianopaz]

Before

with pm.Model() as m:
    a = pm.TruncatedNormal("a", 0, 1, lower=-1, upper=1)

pm.draw(a, mode="JAX") # Fails
pm.draw(a, mode="NUMBA") # Fails

After

with pm.Model() as m:
    a = pm.TruncatedNormal("a", 0, 1, lower=-1, upper=1)

pm.draw(a, mode="JAX") # Works
pm.draw(a, mode="NUMBA") # Works

Context for the issue:

The TruncatedNormal distribution creates a TruncatedNormalRV Op that doesn't have dispatch rules for either JAX or NUMBA. The Truncated class itself seems to work though. In general, it would be nice to know where to write dispatches for these special pymc Ops.

[lucianopaz]

This might be related to #7348

[ricardoV94]

it would be nice to know where to write dispatches for these special pymc Ops

distributions/dispatch/jax
distributions/dispatch/numba

?

[HarshvirSandhu]

I can work on this. Also, would the dispatch for jax go to pymc.sampling.jax?

[ricardoV94]

We may need a dispatch/numba.py and dispatch/jax.py? sampling/jax is too specific

It will need to be based on try except import of jax/numba as those are optional dependecies