code: def moments_of_masked_time_series(time_series_tensor, broadcast_mask):
      """Compute mean and variance, accounting for a mask.
  
      Args:
        time_series_tensor: float `Tensor` time series of shape
          `concat([batch_shape, [num_timesteps]])`.
        broadcast_mask: bool `Tensor` of the same shape as `time_series`.
      Returns:
        mean: float `Tensor` of shape `batch_shape`.
        variance: float `Tensor` of shape `batch_shape`.
      """
      num_unmasked_entries = ps.cast(
          ps.reduce_sum(ps.cast(~broadcast_mask, np.int32), axis=-1),
          time_series_tensor.dtype,
      )
  
      # Manually compute mean and variance, excluding masked entries.
      mean = (
          tf.reduce_sum(
              tf.where(
                  broadcast_mask,
                  tf.zeros([], dtype=time_series_tensor.dtype),
                  time_series_tensor,
              ),
              axis=-1,
          )
          / num_unmasked_entries
      )
      variance = (
          tf.reduce_sum(
              tf.where(
                  broadcast_mask,
                  tf.zeros([], dtype=time_series_tensor.dtype),
                  (time_series_tensor - mean[..., tf.newaxis]) ** 2,
              ),
              axis=-1,
          )
          / num_unmasked_entries
      )
      return mean, variance
  
file_path: probability/tensorflow_probability/python/sts/internal/missing_values_util.py
imports: import collections
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensorshape_util
label_desc: Use image node_id 1 for calling a global function with example usage: moments_of_masked_time_series(time_series_tensor, broadcast_mask) and returns: mean,  variance