Fix bug in storage_input alignment of the JAX backend

[ricardoV94]

Description

When transpiling functions with Shared RNGs to the JAX backend, we replace the original variables by copies so that we can update them in the same container without changing the data type (JAX PRNGs vs numpy Generator), so that the same shared variable can be reused in other functions, which may use distinct backends.

There was a bug in that this replacement could change the position of the input variable in the FunctionGraph, and disalign it with the pre-defined input_storage list. This PR fixes the bug by forcing the position of the new input variable in the FunctionGraph to be the same as the one in input_storage list.

AFAICT the order of the FunctionGraph inputs is not stored/referenced anywhere internally, so this should be safe.

This showed up when trying to use PyMC VI with the JAX backend.

Related Issue

• Related to Allow explicit RNG and Sparse input types in JAX functions #278

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):

[ricardoV94]

CC @abdalazizrashid

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (c5b96d9) 80.92% compared to head (4f79fd4) 80.93%.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #587   +/-   ##
=======================================
  Coverage   80.92%   80.93%           
=======================================
  Files         162      162           
  Lines       46641    46644    +3     
  Branches    11399    11399           
=======================================
+ Hits        37746    37749    +3     
  Misses       6667     6667           
  Partials     2228     2228           

Files	Coverage Δ
pytensor/link/jax/linker.py	95.23% <100.00%> (+0.36%)	⬆️