import liesel.model as lsl
import jax.numpy as jnp
import tensorflow_probability.substrates.jax.distributions as tfd

# create a model
x = lsl.Var(jnp.arange(4).astype(float), name="x")
y = lsl.Var(jnp.arange(4).astype(float), name="y")
batch_index = lsl.Var(
    jnp.array([0, 1]),
    name="batch_index",
)
y_sel_dist = lsl.Dist(lambda x, idx: tfd.Normal(x[idx], 1.0), x, batch_index)
y_sel_val = lsl.Calc(lambda y, idx: y[idx], y, batch_index)
y_sel_val.value
y_sel = lsl.Var(
    y_sel_val,
    y_sel_dist,
    "y_sel",
)

model = lsl.GraphBuilder().add(y_sel_dist).build_model()
lsl.plot_vars(model)