Add JAX support for pt.tri

[jessegrabowski]

Motivation for these changes

pm.LJKCholesky users pt.tri to unpack the cholesky decomposed covariance matrix, so this PR adds support for the underlying Tri Op.

Implementation details

The only wrinkle is the requirement for static shape args -- jnp.tri is exactly the same as jnp.arange in this respect. So I copied the code from the pt.arange jax implementation.

Checklist

• Explain motivation and implementation 👆
• Make sure that the pre-commit linting/style checks pass.
• Link relevant issues, preferably in nice commit messages.
• The commits correspond to relevant logical changes. Note that if they don't, we will rewrite/rebase/squash the git history before merging.
• Are the changes covered by tests and docstrings?
• Fill out the short summary sections 👇

Major / Breaking Changes

None

New features

First step towards support for LKJCorr in PyMC JAX mode

[codecov-commenter]

Codecov Report

Merging #302 (8d3bda6) into main (6b43b43) will decrease coverage by 0.01%.
The diff coverage is 58.82%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #302      +/-   ##
==========================================
- Coverage   80.50%   80.49%   -0.01%     
==========================================
  Files         154      154              
  Lines       45186    45201      +15     
  Branches    11046    11049       +3     
==========================================
+ Hits        36376    36386      +10     
- Misses       6608     6612       +4     
- Partials     2202     2203       +1     

Impacted Files	Coverage Δ
pytensor/link/jax/dispatch/tensor_basic.py	82.35% <58.82%> (-4.19%)	⬇️

... and 1 file with indirect coverage changes

[jessegrabowski]

Tried to re-write the function along the lines you suggested, but it might be too cute. Let me know.

[ricardoV94]

Looks fine, do you want to keep the arange error message changes or revert them?

[jessegrabowski]

Good point, I'll revert it.