How to model structured parameters in Tensorflow Probability distributions?

Question

I'd like to model a multivariate distribution with structured parameters, for example: a multivariate normal with a covariance matrix made up of a low-rank part and a diagonal part. What is the recommended way to achieve this? (Tensorflow 2.8)

DIM=4
mean = tf.Variable(np.zeros(DIM), dtype=tf.float32, name='mean')
low_rank = tf.Variable(np.zeros((DIM, 2)), dtype=tf.float32, name='cov')
diagonal = tf.Variable(np.zeros(DIM), dtype=tf.float32, name='noise')

target_distribution = tfd.MultivariateNormalTriL(
     loc=mean,
     scale_tril=tf.linalg.cholesky(
         tf.linalg.matmul(low_rank, low_rank, transpose_b=True) + tf.linalg.diag(tf.math.softplus(diagonal))
     )
)

print(target_distribution.trainable_variables)

returns only (<tf.Variable 'mean:0' shape=(4,) dtype=float32, numpy=array([0., 0., 0., 0.], dtype=float32)>,), i.e. only those variables assigned directly enter the realm of tracked variables, not those that enter via an expression.

What is the syntax so low_rank and diagonal become trainable variables that I can fit to data?

I realize that there is tfd.MultivariateNormalDiagPlusLowRank that solves this specific example, but I'm still interested in the recommended way to model structured parameters.