def joint_log_prob(self, observed_time_series):
with tf.compat.v1.name_scope(<EOL>'<STR_LIT>', values=[observed_time_series]):<EOL><INDENT>[<EOL>observed_time_series,<EOL>mask<EOL>] = sts_util.canonicalize_observed_time_series_with_mask(<EOL>observed_time_series)<EOL>num_timesteps = distribution_util.prefer_static_value(<EOL>tf.shape(input=observed_time_series))[-<NUM_LIT:2>]<EOL>def log_joint_fn(*param_vals):<EOL><INDENT>"""<STR_LIT>"""<EOL>param_lp = sum([<EOL>param.prior.log_prob(param_val)<EOL>for (param, param_val) in zip(self.parameters, param_vals)<EOL>])<EOL>lgssm = self.make_state_space_model(<EOL>param_vals=param_vals, num_timesteps=num_timesteps)<EOL>observation_lp = lgssm.log_prob(observed_time_series, mask=mask)<EOL>sample_ndims = tf.maximum(<NUM_LIT:0>,<EOL>tf.rank(observation_lp) - tf.rank(param_lp))<EOL>observation_lp = tf.reduce_sum(<EOL>input_tensor=observation_lp, axis=tf.range(sample_ndims))<EOL>return param_lp + observation_lp<EOL><DEDENT><DEDENT>return log_joint_fn<EOL>
Build the joint density `log p(params) + log p(y|params)` as a callable.

        Args:
          observed_time_series: Observed `Tensor` trajectories of shape
            `sample_shape + batch_shape + [num_timesteps, 1]` (the trailing
            `1` dimension is optional if `num_timesteps > 1`), where
            `batch_shape` should match `self.batch_shape` (the broadcast batch
            shape of all priors on parameters for this structural time series
            model). May optionally be an instance of `tfp.sts.MaskedTimeSeries`,
            which includes a mask `Tensor` to specify timesteps with missing
            observations.

        Returns:
         log_joint_fn: A function taking a `Tensor` argument for each model
           parameter, in canonical order, and returning a `Tensor` log probability
           of shape `batch_shape`. Note that, *unlike* `tfp.Distributions`
           `log_prob` methods, the `log_joint` sums over the `sample_shape` from y,
           so that `sample_shape` does not appear in the output log_prob. This
           corresponds to viewing multiple samples in `y` as iid observations from a
           single model, which is typically the desired behavior for parameter
           inference.