Composition of Joint Distributions in TensorFlow Probability

I am new to probabilistic programming and TFP. I am trying to use it for probabilistic graphical modelling, which appeals to me for its flexibility in breaking down complex models into simpler piecewise explanations. However I am having difficulty in defining higher level models which are probabilistically related to variables defined by lower level models.

Suppose I define the following distribution for a random variable theta_1 based on explanatory variables x_1, x_2 and likewise for theta_2 based on x_3, x_4:

theta_1_model = tfd.JointDistributionSequential(
    [
        tfd.Normal(0, 1, name='x_1'),
        tfd.Normal(0, 1, name='x_2'),
        lambda x_2, x_1: tfd.Normal(x_1, tf.exp(x_2), name='theta_1')
    ]
)
theta_2_model = tfd.JointDistributionSequential(
    [
        tfd.Normal(0, 1, name='x_3'),
        tfd.Normal(0, 1, name='x_4'),
        lambda x_4, x_3: tfd.Normal(x_3, tf.exp(x_4), name='theta_2')
    ]
)

I would then like to define models which rely on theta_1 and theta_2 as explanatory variables. For example suppose I would like to model a random variable y which is normally distributed about the difference between theta_1 and theta_2. Given that theta_1_model and theta_2_model are themselves distributions (for dim-3 random variables) it seems like I should be able to use them within a JointDistributionSequential definition for y along the lines of:

y_model = tfd.JointDistributionSequential(
    [
        theta_1_model,
        theta_2_model,
        lambda t_2_m, t_1_m: tfd.Normal(t_1_m[...,2]-t_2_m[...,2], 1, name='y')
    ]
)

where here I have attempted to use [..., 2] indexing to reference the theta component within the dim-3 distributions. This does not work however. It is in specifying the relationship between new variables in the higher-level model and individual variables already defined in the lower level models that I am having difficulty.

I could of course define the full model for y in one big JointDistributionSequential definition. However for more complicated applications where I'd like to reuse the lower level components this is infeasible.

There has recently been a white-paper (https://arxiv.org/pdf/2001.11819.pdf) published entitled 'Joint Distributions for TensorFlow Probability' with a 'Composition of Joint Distributions' section. This section suggests that joint distributions can in principle be combined in the way I'm suggesting, but doesn't show how the relationships between new variables in the higher-level model and variables already defined in the lower level models may be specified.

Any help would be greatly appreciated. Thanks.