Error with jit and batched sampling

The following code samples three times from the same model. Only the final attempt, which combines jax.jit and sample_shape, raises a ValueError. This is on today's tfp-nightly (0.12.0-dev20201030) and jax from earlier today (21bb501) on CPU

import jax
from tensorflow_probability.substrates import jax as tfp

tfd = tfp.distributions
Root = tfd.JointDistributionCoroutineAutoBatched.Root

@tfd.JointDistributionCoroutineAutoBatched
def model():
    x = yield Root(tfd.Sample(
        tfd.Uniform(0, 1),
        sample_shape=(),
    ))
    return x

def sample():
    return model.sample(
        sample_shape=(1,), 
        seed=jax.random.PRNGKey(0)
    )
# Works as expected without jit
sample()

@jax.jit
def sample():
    return model.sample(
        seed=jax.random.PRNGKey(0)
    )
# Works as expected with default sample shape
sample()

@jax.jit
def sample():
    return model.sample(
        sample_shape=(1,), 
        seed=jax.random.PRNGKey(0)
    )
# Raises value error now that jit and sample_shape are used
sample()

> ValueError: vmap must have at least one non-None value in in_axes

stack trace

Start by running the supplied reproduction with the reported tfp-nightly and JAX versions, then follow the linked stack trace from model.sample into the JAX sampling path. Done means the jitted call with sample_shape=(1,) no longer raises the reported ValueError and the behavior is covered by regression testing.