Splitting before using JAX key

[AdrienCorenflos]

This closes #1344
As discussed offline with @rlouf at the moment the jax splitting logic relies on knowing the internals on the splitting and doesn't follow the JAX best practices. This PR fixes this.

Here are a few important guidelines and requirements to check before your PR can be merged:

• There is an informative high-level description of the changes.
• The description and/or commit message(s) references the relevant GitHub issue(s).
• pre-commit is installed and set up.
• The commit messages follow these guidelines.
• The commits correspond to relevant logical changes, and there are no commits that fix changes introduced by other commits in the same branch/BR.
• There are tests covering the changes introduced in the PR.

[AdrienCorenflos]

I ran pre-commit, should be good

[AdrienCorenflos]

@rlouf

[rlouf]

Should be good to merge once the tests pass. Thanks for spotting this and fixing it!

[codecov]

Codecov Report

Merging #1345 (16b50ba) into main (2434cb4) will increase coverage by 0.00%.
The diff coverage is 100.00%.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1345   +/-   ##
=======================================
  Coverage   74.35%   74.35%           
=======================================
  Files         177      177           
  Lines       49046    49056   +10     
  Branches    10379    10379           
=======================================
+ Hits        36468    36478   +10     
  Misses      10285    10285           
  Partials     2293     2293           

Impacted Files	Coverage Δ
aesara/link/jax/dispatch/random.py	100.00% <100.00%> (ø)