In JAX random linking splitting should happen before using the key

[AdrienCorenflos]

See for example

aesara/aesara/link/jax/dispatch/random.py

Lines 126 to 130 in fe3e76d

	def sample_fn(rng, size, dtype, *parameters):
	rng_key = rng["jax_state"]
	sample = jax_op(rng_key, *parameters, shape=size, dtype=dtype)
	rng["jax_state"] = jax.random.split(rng_key, num=1)[0]
	return (rng, sample)

in order to be internals-agnostic, this should be

def sample_fn(rng, size, dtype, *parameters): 
     rng_key = rng["jax_state"] 
     rng_key, sample_key = jax.random.split(rng_key, 2) 
     sample = jax_op(sample_key, *parameters, shape=size, dtype=dtype) 
     rng["jax_state"] = rng_key
     return (rng, sample)