Как обрабатывать разделение PRNG в контексте jax.vmap? ⇐ Python

[Anonymous]

У меня есть функция, которая моделирует стохастическое дифференциальное уравнение. В настоящее время, без стохастического шума, мой вызов моделирования процесса до момента времени t выглядит следующим образом (и да, мне нужно использовать jax):

Код: Выделить всё

def evolve(u, t):
# return u + dt * b(t, u) + sigma(t, u) * sqrt_dt * noise

def simulate(x, t):
k = jax.numpy.floor(t / dt).astype(int)
u = jax.lax.fori_loop(0, k, lambda i, u : evolve(u, i * dt), u)

Теперь вместе с шумом приходит боль. Я знаю C++, и мне лишь изредка приходится использовать Python для исследований/научной работы. И я действительно не понимаю, как мне нужно (или следует) реализовать здесь разделение PRNG. Думаю, я бы изменил на

Код: Выделить всё

def evolve(u, t, key):
noise = jax.random.multivariate_normal(key, jax.numpy.zeros(d), covariance_matrix, shape = (n,))
# return u + dt * b(t, u) + sigma(t, u) * sqrt_dt * noise

Но я думаю, это не будет работать должным образом. Если я все понял правильно, мне нужно использовать jax.random.split, чтобы разделить ключ. Потому что, если я этого не сделаю, я получу коррелирующие образцы. Но как и где мне нужно разделить?
А еще: я думаю, мне нужно будет изменить имитацию на def моделирования(x, t, key). Но тогда должна ли имитация также возвращать измененный ключ?
И чтобы еще больше усложнить задачу: я фактически обертываю имитацию в пакетный_симулят< /code> функция, которая использует jax.vmap для обработки целого пакета символов x и t. Как передать PRNG в эту функцию пакетного_симуляции, как передать его (и передать его) в jax.vmap и что должен вернуть пакетный_форвард? На первый взгляд мне кажется, что нужно взять один PRNG и разбить его на множество (из-за vmap). Но что тогда делает вызывающая сторона пакетного_форварда...
Полностью потерялся в этом. Любая помощь приветствуется!

Подробнее здесь: https://stackoverflow.com/questions/792 ... ap-context