code: def initial_value_of_masked_time_series(
      time_series_tensor, broadcast_mask
  ):
      """Get the first unmasked entry of each time series in the batch.
  
      Args:
        time_series_tensor: float `Tensor` of shape `batch_shape + [num_timesteps]`.
        broadcast_mask: bool `Tensor` of same shape as `time_series`.
      Returns:
        initial_values: float `Tensor` of shape `batch_shape`.  If a batch element
          has no unmasked entries, the corresponding return value for that element
          is undefined.
        first_unmasked_indices: int `Tensor` of shape `batch_shape` -- the index of
          the first unmasked entry for each time series in the batch.  If there are
          no unmasked entries, the returned index is the length of the series.
      """
  
      num_timesteps = ps.shape(time_series_tensor)[-1]
  
      # Compute the index of the first unmasked entry for each series in the batch.
      unmasked_negindices = ps.cast(
          ~broadcast_mask, np.int32
      ) * ps.range(num_timesteps, 0, -1)
      first_unmasked_indices = num_timesteps - ps.reduce_max(
          unmasked_negindices, axis=-1
      )
      # Avoid out-of-bounds errors if all indices are masked.
      safe_unmasked_indices = ps.minimum(
          first_unmasked_indices, num_timesteps - 1
      )
  
      batch_dims = tensorshape_util.rank(safe_unmasked_indices.shape)
      if batch_dims is None:
          raise NotImplementedError(
              "Cannot compute initial values of a masked time series with"
              "dynamic rank."
          )  # `batch_gather` requires static rank
  
      # Extract the initial value for each series in the batch.
      initial_values = tf.squeeze(
          tf.gather(
              params=time_series_tensor,
              indices=safe_unmasked_indices[..., np.newaxis],
              batch_dims=batch_dims,
              axis=-1,
          ),
          # Since we've gathered exactly one step from the `num_timesteps` axis,
          # we can remove that axis entirely.
          axis=-1,
      )
  
      return initial_values, first_unmasked_indices
  
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 2 for calling a global function with example usage: initial_value_of_masked_time_series(time_series_tensor, broadcast_mask) and returns: initial_values,  first_unmasked_indices