Address some Jax random key issues

[codecov[bot]]

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

```diff @@ 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](https://app.codecov.io/gh/theislab/pertpy/pull/531?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=theislab) | Coverage Δ | | |---|---|---| | [pertpy/metadata/\_cell\_line.py](https://app.codecov.io/gh/theislab/pertpy/pull/531?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=theislab#diff-cGVydHB5L21ldGFkYXRhL19jZWxsX2xpbmUucHk=) | `62.87% <100.00%> (ø)` | | | [pertpy/metadata/\_moa.py](https://app.codecov.io/gh/theislab/pertpy/pull/531?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=theislab#diff-cGVydHB5L21ldGFkYXRhL19tb2EucHk=) | `79.48% <100.00%> (ø)` | | | [pertpy/plot/\_augur.py](https://app.codecov.io/gh/theislab/pertpy/pull/531?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=theislab#diff-cGVydHB5L3Bsb3QvX2F1Z3VyLnB5) | `44.11% <ø> (ø)` | | | [pertpy/plot/\_coda.py](https://app.codecov.io/gh/theislab/pertpy/pull/531?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=theislab#diff-cGVydHB5L3Bsb3QvX2NvZGEucHk=) | `100.00% <ø> (ø)` | | | [pertpy/plot/\_guide\_rna.py](https://app.codecov.io/gh/theislab/pertpy/pull/531?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=theislab#diff-cGVydHB5L3Bsb3QvX2d1aWRlX3JuYS5weQ==) | `53.84% <ø> (ø)` | | | [pertpy/plot/\_milopy.py](https://app.codecov.io/gh/theislab/pertpy/pull/531?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=theislab#diff-cGVydHB5L3Bsb3QvX21pbG9weS5weQ==) | `50.00% <ø> (ø)` | | | [pertpy/plot/\_mixscape.py](https://app.codecov.io/gh/theislab/pertpy/pull/531?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=theislab#diff-cGVydHB5L3Bsb3QvX21peHNjYXBlLnB5) | `76.92% <ø> (ø)` | | | [pertpy/preprocessing/\_guide\_rna.py](https://app.codecov.io/gh/theislab/pertpy/pull/531?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=theislab#diff-cGVydHB5L3ByZXByb2Nlc3NpbmcvX2d1aWRlX3JuYS5weQ==) | `54.76% <ø> (ø)` | | | [pertpy/tools/\_coda/\_sccoda.py](https://app.codecov.io/gh/theislab/pertpy/pull/531?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=theislab#diff-cGVydHB5L3Rvb2xzL19jb2RhL19zY2NvZGEucHk=) | `90.00% <100.00%> (ø)` | | | [pertpy/tools/\_dialogue.py](https://app.codecov.io/gh/theislab/pertpy/pull/531?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=theislab#diff-cGVydHB5L3Rvb2xzL19kaWFsb2d1ZS5weQ==) | `13.70% <ø> (ø)` | | | ... and [12 more](https://app.codecov.io/gh/theislab/pertpy/pull/531?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=theislab) | |