Sampling from BCOO containing multiple categorical distributions.

[tomsch420]

I have a usecase where I have the log probabilities of many categorical distributions in a BCOO array. Now I want samples from each row with a different quantity.

The method that does the required operation is here:

def sample_from_sparse_probabilities(log_probabilities: BCOO, amount: jax.Array, key: jax.random.PRNGKey) -> BCOO:
    """
    Sample from a sparse array of probabilities.
    Each row in the sparse array encodes a categorical probability distribution.

    :param log_probabilities: The unnormalized sparse array of log-probabilities.
    :param amount:  The amount of samples to draw from each row.
    :param key: The random key.
    :return: The samples that are drawn for each state in the probabilities indicies.
    """
    all_samples = []

    for probability_row, row_amount in zip(log_probabilities, amount):
        probability_row: BCOO
        probability_row = probability_row.sum_duplicates(remove_zeros=False)

        samples = jax.random.categorical(key, probability_row.data, shape=(row_amount.item(), ))
        frequencies = jnp.zeros((probability_row.data.shape[0],), dtype=jnp.int32)
        frequencies = frequencies.at[samples].add(1)
        all_samples.append(frequencies)

    return BCOO((jnp.concatenate(all_samples), log_probabilities.indices), shape=log_probabilities.shape,
                indices_sorted=True, unique_indices=True)

A minimal code block that runs this is

probs = BCOO.fromdense(jnp.array([[0.1, 0.2, 0., .7],
                                [0.4, 0., 0.6, 0.]]))
probs.data = jnp.log(probs.data)
amount = jnp.array([2, 3])

samples = sample_from_sparse_probabilities(probs,amount, jax.random.PRNGKey(69))

Unfortunatly, this is very slow. The perfetto trace for instance states that the entire genexpr thing is taken 0.5s and that sum_duplicates alone is also taking 0.5s. Does anyone have an idea how to increase the speed of such an operation? I already tried vmap and jit, but they cannot wrap such a function, nor the slow parts of that function.

[jakevdp]

Unfortunately, BCOO is very slow: XLA is a non-sparse-aware compiler, and BCOO is a reference implementation of sparse operations within that dense compiler context; it was never designed to be particularly performant. If it's useful for you, you should use it. If it's not useful, you shouldn't.