Add `MultinomialRV` JAX implementation #1360

GStechschulte · 2022-12-11T16:20:24Z

This draft PR is a work in progress and contains a JAX implementation of MultinomialRV for issue #1326. The implementation builds off the Multinomial Distribution implementation in NumPyro. Likewise, the output is similar to that of the numpy implementation. Below, you will find a brief outline of the functions used to construct the MultinomialRV.

def _categorical(key, p, shape)

returns the outcomes $k$ with probability $p$ for each trial / experiment $n$.

def _scatter_add_ones(operand, indices, update)

returns the outcome counts by utilising the jax.lax.scatter_add() function
operand is a zero filled array.
indices is the outcomes array with an added dimension and specifies the indices to which the update should be applied to.
update is an array filled with ones and can be thought of as a cnt += 1 for each $K = k$ occurrence.
In summary, the operand array is updated +1 using the update array according to the outcomes in the indices array.

I still need to add a test for this. Thanks!

rlouf · 2022-12-12T09:26:52Z

Thanks! We also need to figure out if the licenses are compatible and how to do proper attribution if you took inspiration from someone else's implementation. It looks like Numpyro is licensed under Apache 2.0

GStechschulte · 2022-12-12T20:51:22Z

Thanks! We also need to figure out if the licenses are compatible and how to do proper attribution if you took inspiration from someone else's implementation. It looks like Numpyro is licensed under Apache 2.0

Based on the NumPyro Apache License 2.0 section 4, we may reproduce and distribute copies of the Work or Derivative Works (the JAX implementation of MultinomialRV) provided we:

give any other recipients of the Work or Derivative Works a copy of this License; and
the modified file must contain carry a notice stating the file was changed
in the Source form of any Derivative Works that You distribute, all copyright, patent, trademark, and attribution notices from the Source form of the Work
if the Work includes a "NOTICE" text file as part of its distribution, then any Derivative Works that You distribute must include a readable copy of the attribution notices contained within such NOTICE file, excluding those notices that do not pertain to any part of the Derivative Works, in at least one of the following places: within a NOTICE text file distributed as part of the Derivative Works; within the Source form or documentation

Since inspiration was drawn from a few functions and not an entire file, I suggest we include, in addition to (1), (2), (3), in the documentation for this RV, something along the lines

MultinomialRV uses source code from the file xyz.py from of the NumPyro project, copyright YYYY, licensed under the Apache 2.0 license>

rlouf · 2022-12-13T13:12:19Z

aesara/link/jax/dispatch/random.py

+    def _categorical(key, p, shape):
+        shape = shape or p.shape[:-1]
+        s = jax.numpy.cumsum(p, axis=-1)
+        r = jax.random.uniform(key, shape=shape + (1,))
+
+        return jax.numpy.sum(s < r, axis=-1)


I am surprised because JAX does have an implementation for the categorical distribution here that uses their implementation for the Gumbel distribution. Is this justified in the codebase, or is it just because it was implemented at a time where jax.random.categorical was not available (which you should be able to determine with git blame)?

rlouf · 2022-12-13T13:20:01Z

aesara/link/jax/dispatch/random.py

+        samples_2d = jax.vmap(_scatter_add_one, (0, 0, 0))(
+            jax.numpy.zeros((outcomes_2d.shape[0], p.shape[-1]), dtype=outcomes.dtype),
+            jax.numpy.expand_dims(outcomes_2d, axis=-1),
+            jax.numpy.ones(outcomes_2d.shape, dtype=outcomes.dtype)
+            )
+
+        sample = jax.numpy.reshape(samples_2d, size + p.shape[-1:])


Couldn't we use jax.nn.one_hot on the output of the categorical and then reduce the resulting tensor?

Yeah, your proposal seems to be a more elegant solution; thanks. Just committed.

rlouf · 2022-12-14T13:32:29Z

I rebased your branch on main to use the new key splitting scheme in the JAX backend. You'll have to pull the changes!

GStechschulte

When I performed the git pull --merge, it kept a copy of the previous code (without the new key splitting). This was likely a mistake on my end. Therefore, I deleted it.

GStechschulte · 2022-12-13T19:17:39Z

aesara/link/jax/dispatch/random.py

+        samples_2d = jax.vmap(_scatter_add_one, (0, 0, 0))(
+            jax.numpy.zeros((outcomes_2d.shape[0], p.shape[-1]), dtype=outcomes.dtype),
+            jax.numpy.expand_dims(outcomes_2d, axis=-1),
+            jax.numpy.ones(outcomes_2d.shape, dtype=outcomes.dtype)
+            )
+
+        sample = jax.numpy.reshape(samples_2d, size + p.shape[-1:])


Yeah, your proposal seems to be a more elegant solution; thanks. Just committed.

rlouf · 2022-12-14T21:17:37Z

Yes you need to git pull --rebase in such cases. Here's a good explanation of how rebasing works. And of course the documentation for git pull

rlouf · 2022-12-15T16:00:05Z

aesara/link/jax/dispatch/random.py

+    def _categorical(key, p, shape):
+        shape = shape or p.shape[:-1]
+        s = jax.numpy.cumsum(p, axis=-1)
+        r = jax.random.uniform(key, shape=shape + (1,))
+
+        return jax.numpy.sum(s < r, axis=-1)


Is that different from jax.random.categorical?

Yes, the jax.random.categorical uses their implementation of the Gumbel distribution. However, after using git blame, it seems the NumPyro team used this implementation before the implementation of the jax.random.categorical.

If we would like to update this code to use the jax.random.categorical, it would be the following:

def sample_fn(rng, size, dtype, *parameters): rng_key = rng["jax_state"] rng_key, sampling_key = jax.random.split(rng_key, 2) n, p = parameters n_max = jax.numpy.max(n) size = size or p.shape[:-1] logits = jax.scipy.special.logit(p) indices = jax.random.categorical(jax_key, logits, shape=(n_max,) + size) one_hot = jax.nn.one_hot(indices, p.shape[0]) sample = jax.numpy.sum(one_hot, axis=0, dtype=dtype, keepdims=False) rng["jax_state"] = rng_key return (rng, sample)

That looks much simpler, great!

Since the multinomial distribution is slightly more complex than the other ones when it comes to shapes we should make sure that the output shape of the samples that are generated in the JAX backend is identical to those of the samples generated with the other backends.

I tend to use jax.random.choice for this kind of thing. jax.random.categorical (via Gumbel) has a quadratic complexity.

You just need to check the Gumbel implementation, they form a N by N array.

brandonwillard · 2023-01-01T19:48:08Z

.gitignore

@@ -55,3 +55,4 @@ aesara-venv/
 testing-report.html
 coverage.xml
 .coverage.*
+jax_multinomial_test.py


This looks like a job for a "local" Git ignore (see here).

rlouf · 2023-01-16T19:53:38Z

@GStechschulte I think we should follow @AdrienCorenflos's suggestion here. I'll take another look at it this week.

rlouf · 2023-02-21T15:23:53Z

@GStechschulte I rebased your branch on main. Do you plan on implementing @AdrienCorenflos's suggestion above?

brandonwillard

I've squared, rebased, and added a few fixes. The JAX steps are not properly accounting for the shapes of the distribution parameters, so that needs to be finished. I refactored the tests so that they cover one of the most basic cases and another that should confirm that the sizes/shapes are handled correctly (when they are).

@rlouf, I had to add a case for Constants in assert_size_argument_jax_compatible. You'll need to confirm that this is valid more generally.

FYI: the tests aren't being run with shape inference, since the jax_mode used by the tests doesn't include "ShapeOpt" (or whatever its tag is). Standard JAX mode should, since it's included with the "fast_run" tag, but, if we want to test for non-trivial shape scenarios (e.g. the shape value isn't explicitly constant, but can be "inferred" as a constant value), we'll need to add it.

rlouf · 2023-03-10T14:08:53Z

@rlouf, I had to add a case for Constants in assert_size_argument_jax_compatible. You'll need to confirm that this is valid more generally.

That's valid. I was so focused on the complex case that I forgot the simplest one.

brandonwillard added JAX Involves JAX transpilation random variables Involves random variables and/or sampling labels Dec 11, 2022

This was linked to issues Dec 12, 2022

Add JAX implementation for MultinomialRV #1326

Open

Add JAX implementation for HypergeometricRV #1324

Open

rlouf removed a link to an issue Dec 12, 2022

Add JAX implementation for HypergeometricRV #1324

Open

rlouf changed the title ~~Add JAX implementation of MultinomialRV~~ Add JAX implementation MultinomialRV Dec 12, 2022

rlouf changed the title ~~Add JAX implementation MultinomialRV~~ Add MultinomialRV JAX implementation Dec 12, 2022

GStechschulte marked this pull request as ready for review December 12, 2022 21:00

rlouf reviewed Dec 13, 2022

View reviewed changes

rlouf force-pushed the jax_multinomial branch from b83b1ac to 3b006eb Compare December 14, 2022 13:26

GStechschulte commented Dec 14, 2022

View reviewed changes

rlouf reviewed Dec 15, 2022

View reviewed changes

brandonwillard reviewed Jan 1, 2023

View reviewed changes

rlouf force-pushed the jax_multinomial branch from 8932215 to 662c586 Compare February 21, 2023 15:22

brandonwillard force-pushed the jax_multinomial branch from 662c586 to d93fb8d Compare March 8, 2023 00:49

brandonwillard and others added 2 commits March 9, 2023 20:02

Allow constant size arguments during JAX transpilation

3bd2559

Add a JAX implementation of MultinomialRV

148217a

brandonwillard force-pushed the jax_multinomial branch from d93fb8d to 148217a Compare March 10, 2023 02:02

brandonwillard reviewed Mar 10, 2023

View reviewed changes

GStechschulte closed this by deleting the head repository Mar 12, 2024

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add `MultinomialRV` JAX implementation #1360

Add `MultinomialRV` JAX implementation #1360

GStechschulte commented Dec 11, 2022

rlouf commented Dec 12, 2022 •

edited

GStechschulte commented Dec 12, 2022 •

edited

rlouf Dec 13, 2022

rlouf Dec 13, 2022

GStechschulte Dec 13, 2022

rlouf commented Dec 14, 2022

GStechschulte left a comment

GStechschulte Dec 13, 2022

rlouf commented Dec 14, 2022 •

edited

rlouf Dec 15, 2022

GStechschulte Dec 15, 2022

rlouf Dec 16, 2022

AdrienCorenflos Dec 22, 2022 •

edited

AdrienCorenflos Dec 22, 2022

brandonwillard Jan 1, 2023

rlouf commented Jan 16, 2023

rlouf commented Feb 21, 2023

brandonwillard left a comment

rlouf commented Mar 10, 2023

Add MultinomialRV JAX implementation #1360

Add MultinomialRV JAX implementation #1360

Conversation

GStechschulte commented Dec 11, 2022

rlouf commented Dec 12, 2022 • edited

GStechschulte commented Dec 12, 2022 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rlouf commented Dec 14, 2022

GStechschulte left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rlouf commented Dec 14, 2022 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

AdrienCorenflos Dec 22, 2022 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rlouf commented Jan 16, 2023

rlouf commented Feb 21, 2023

brandonwillard left a comment

Choose a reason for hiding this comment

rlouf commented Mar 10, 2023

Add `MultinomialRV` JAX implementation #1360

Add `MultinomialRV` JAX implementation #1360

rlouf commented Dec 12, 2022 •

edited

GStechschulte commented Dec 12, 2022 •

edited

rlouf commented Dec 14, 2022 •

edited

AdrienCorenflos Dec 22, 2022 •

edited