Add support for sample_shape to primitive distributions (#1576)

This PR passes any `sample_shape` argument provided to a tfp
distribution through to the `sample` method, instead of passing it to
the constructor.

@femtomc one issue that I hit here was that to use this, I have to wrap
the argument in `Const`, because it seems like our code tries to trace
non-jit-compiled fns, or something like that. I'll add the error I saw
as a reply.

Co-authored-by: Mathieu Huot <[email protected]>

Author: sritchie

src/genjax/_src/core/pytree.py

@@ -250,7 +250,7 @@ class Const(Generic[R], Pytree):
 
 
         def f(c):
-            if c.const == 5:
+            if c.unwrap() == 5:
                 return 10.0
             else:
                 return 5.0

src/genjax/_src/generative_functions/distributions/tensorflow_probability/__init__.py

@@ -16,6 +16,7 @@
 import jax.numpy as jnp
 from tensorflow_probability.substrates import jax as tfp
 
+from genjax._src.core.pytree import Const
 from genjax._src.core.typing import Array, Callable
 from genjax._src.generative_functions.distributions.distribution import (
     ExactDensity,
@@ -49,11 +50,15 @@ def tfp_distribution(
     """
 
     def sampler(key, *args, **kwargs):
+        sample_shape = kwargs.pop("sample_shape", ())
         d = dist(*args, **kwargs)
-        return d.sample(seed=key)
+        return d.sample(seed=key, sample_shape=Const.unwrap(sample_shape))
 
     def logpdf(v, *args, **kwargs):
+        # Remove unused kwarg to match sampler function behavior
+        kwargs.pop("sample_shape", ())
         d = dist(*args, **kwargs)
+
         return d.log_prob(v)
 
     return exact_density(sampler, logpdf, name or dist.__name__)

tests/generative_functions/test_static_gen_fn.py

@@ -84,6 +84,14 @@ def annotated_function(x: float, y: float) -> float:
 
 
 class TestMisc:
+    def test_static_sample_shape(self):
+        @genjax.gen
+        def f():
+            return genjax.normal(0.0, 1.0, sample_shape=genjax.Const((2, 2))) @ "normal"
+
+        tr = f.simulate(jax.random.key(0), ())
+        assert tr.get_retval().shape == (2, 2)
+
     def test_switch_chm_and_static(self):
         @genjax.gen
         def model():