Add jax dispatch for KroneckerProduct Op

[jessegrabowski]

Description

#622 broke the JAX dispatch for pt.nlinalg.kron, because we're wrapping it in a dummy OpFromGraph(inline=False). I tried just changing it to inline=True, but got some dynamic shape errors when the graph was converted to JAX. I think it's easier to just directly register the new OpFromGraph to jnp.kron, which this PR does. It also adds a test.

Related Issue

• Closes #
• Related to Add rewrite to lift linear algebra through certain linalg ops #622

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.87%. Comparing base (0e29d76) to head (e21e50b).
Report is 161 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #822   +/-   ##
=======================================
  Coverage   80.87%   80.87%           
=======================================
  Files         163      163           
  Lines       46847    46852    +5     
  Branches    11463    11464    +1     
=======================================
+ Hits        37887    37892    +5     
  Misses       6750     6750           
  Partials     2210     2210           

Files with missing lines	Coverage Δ
pytensor/link/jax/dispatch/nlinalg.py	91.01% <100.00%> (+0.53%)	⬆️