wektoryzacja rozkładu multinomial w pythonie

Question

Cześć, zastanawiam się, czy istnieje jakiś sposób na wektoryzację funkcji np.random.multinomial albo torch.multinomial w taki sposób by nie przyjmowały liczby prób i wektora prawdopodobieństw, tylko wektor ilości prób i macierz prawdopodobieństw, co sprawi, że każde generowanie będzie posiadać inne parametry (w przeciwieństwie do parametru size, który tylko powiela ten sam eksperyment). 

Kod pokazujący dokładnie o co chodzi:

import numpy as np

# dummy data for testing
tries = np.random.randint(1000, 10000, size=1000)
probs = np.random.rand(1000, 3)
probs /= probs.sum(axis=1, keepdims=True)

# output array
output = np.zeros((1000, 3), dtype=int)

# TODO vectorize this loop 
for i in range(len(tries)):
    output[i, :] = np.random.multinomial(tries[i], probs[i])

Na internecie znalazłem jedynie coś takiego https://stackoverflow.com/questions/36952419/vectorizing-numpy-random-multinomial. Niestety tutaj problem jest zredukowany do jednej próby czyli defacto losowego wyboru z różnymi prawdopodobieństwami dla elementów.

Zastanawiam się czy istnieje jakiś mądry sposób, który niekoniecznie korzysta z funkcji np.random.multinomial, ale pozwala na znaczące przyśpieszenie takiego eksperymentu poprzez pozbycie się pętli.

Answer 1

Znalazłem odpowiedź na to pytanie. Zdecydowałem się odpowiedzieć zamiast usuwać, jeśli komuś się przyda (jakkolwiek nie byłoby to mało prawdopodobne).

tries = np.random.randint(1000, 10000, size=1000)
probs = np.random.rand(1000, 3)
probs /= probs.sum(axis=1, keepdims=True)

# vectorized version
rng = np.random.default_rng()
output = rng.multinomial(tries, probs)

Z pomiarów %timeit wynika, że to rozwiązanie jest ponad 10 razy szybsze.