import tensorflow_probability as tfp
import jax
import jax.numpy as jnp

tfd = tfp.substrates.jax.distributions
tfb = tfp.substrates.jax.bijectors

key = jax.random.PRNGKey(42)

# This works
mix = tfd.Mixture(
    cat=tfd.Categorical(probs=jnp.asarray([0.3, 0.7])),
    components=[tfd.Normal(0.0, 1.0), tfd.Normal(1., 2.)],
)
print("Here it works. This is a sample: ", mix.sample(3, key))


mean = jnp.zeros(5)
covariance_matrix = jnp.eye(5)
dist = tfd.MultivariateNormalFullCovariance(loc=mean, covariance_matrix=covariance_matrix)

split_dist = tfd.TransformedDistribution(
    distribution=dist,
    bijector=tfb.Split((2, 3)),
)

bijectors = [tfb.Exp(), tfb.Sigmoid()]  # Note: has to be a list, for tuple doesn't work

bij = tfb.JointMap(bijectors)
tr_dist = tfd.TransformedDistribution(distribution=split_dist, bijector=bij)

print("There is an error:")
tfd.Mixture(
    cat=tfd.Categorical(probs=jnp.asarray([0.3, 0.7])),
    components=[split_dist, tr_dist],
)