`JointDistributionCoroutineAutoBatched.sample_distributions` errors using jax substrate

Hi, I found that sample_distributions can error when using the jax substrate. sample seems to be ok, using tf appears to be ok, and using the non auto batched joint distribution is also ok.

Here is a small example:

from functools import partial

import jax
from tensorflow_probability.substrates import jax as tfp

tfd = tfp.distributions

tfp.__version__
# 0.25.0

@partial(tfd.JointDistributionCoroutine, batch_ndims=0)
def joint_dist():
    x = yield tfd.Gamma(2.0, 10.0, name="x")
    y = yield tfd.Gamma(x, 10.0, name="y")


seed = jax.random.key(123)

# ok
dists, samples = joint_dist.sample_distributions(x=[1.0, 2.0], seed=seed)

# samples:
# StructTuple(
#   x=Array([1., 2.], dtype=float32),
#   y=Array([0.10665689, 0.21802416], dtype=float32)
# )


@tfd.JointDistributionCoroutineAutoBatched
def joint_dist():
    x = yield tfd.Gamma(2.0, 10.0, name="x")
    y = yield tfd.Gamma(x, 10.0, name="y")


# ok
samples = joint_dist.sample(x=[1.0, 2.0], seed=seed)

# samples:
# StructTuple(
#   x=Array([1., 2.], dtype=float32),
#   y=Array([0.05508393, 0.14792603], dtype=float32)
# )

# ValueError: Attempt to convert a value (<object object at 0x717aa57398a0>) with an unsupported type (<class 'object'>) to a Tensor.
dists, samples = joint_dist.sample_distributions(x=[1.0, 2.0], seed=seed)

Thanks!

Start by running the provided JAX reproduction for JointDistributionCoroutineAutoBatched.sample_distributions and compare it with the working sample and non-auto-batched cases. Trace the sample_distributions entry point and the unsupported object-to-Tensor conversion; done means the auto-batched JAX call returns distributions and samples without the ValueError.