Address some Jax random key issues

[Zethson]

No description provided.

[codecov]

Codecov Report

Attention: 7 lines in your changes are missing coverage. Please review.

Comparison is base (a98cb3d) 64.07% compared to head (88efb18) 64.06%.
Report is 11 commits behind head on main.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #531      +/-   ##
==========================================
- Coverage   64.07%   64.06%   -0.01%     
==========================================
  Files          43       43              
  Lines        4982     4981       -1     
==========================================
- Hits         3192     3191       -1     
  Misses       1790     1790              

Files	Coverage Δ
pertpy/metadata/_cell_line.py	62.87% <100.00%> (ø)
pertpy/metadata/_moa.py	79.48% <100.00%> (ø)
pertpy/plot/_augur.py	44.11% <ø> (ø)
pertpy/plot/_coda.py	100.00% <ø> (ø)
pertpy/plot/_guide_rna.py	53.84% <ø> (ø)
pertpy/plot/_milopy.py	50.00% <ø> (ø)
pertpy/plot/_mixscape.py	76.92% <ø> (ø)
pertpy/preprocessing/_guide_rna.py	54.76% <ø> (ø)
pertpy/tools/_coda/_sccoda.py	90.00% <100.00%> (ø)
pertpy/tools/_dialogue.py	13.70% <ø> (ø)
... and 12 more