Hello, I am currently trying to fit a hidden Markov model on a sequence of observations that consist of angles and speeds. The observations are taken from a dataset in which the amount of time between them is not constant. However, I want to model a sequence of states with a constant amount of time between them. In order to do this, I am trying to model “missing” observations, i.e. states for which an observation does not exist. I am trying to do this with a mixture distribution where the first element of each event indicates whether the observation is missing or not. Missing observations therefore correspond to [0, 0, 0], while non-missing observations correspond to [1, angle, speed], where angle is in [-pi, pi] and speed is in [0, infinity). However, I am getting an error I don’t understand when trying to compute the log probability of the hidden Markov model. Here is my current code:
# assume 2 states
num_states = 2

# Randomly initialize the initial state distribution as well as the transition probabilities
initial_probs = tf.Variable(scipy.special.softmax(rng.random([num_states])), name='initial_probs', dtype=tf.float32)
transition_probs = tf.Variable(scipy.special.softmax(rng.random([num_states, num_states]), axis=1), name='transition_probs', dtype=tf.float32)

# Initialize locations and concentrations of Von Mises distributions for turning angles
vm_locs = tf.Variable(np.zeros(num_states), dtype=tf.float32)
vm_cons = tf.Variable(np.zeros(num_states), dtype=tf.float32)

# Initialize shapes and rates of Gamma distributions for speed
gamma_shapes = tf.Variable(np.ones(num_states), dtype=tf.float32)
gamma_rates = tf.Variable(np.ones(num_states), dtype=tf.float32)

mixed_dists = tfd.Mixture(
    cat=tfd.Categorical(probs=[[0.5, 0.5]] * num_states),
    components=[
        tfd.Independent(tfd.Categorical(probs=[[[1.0, 0.0], [1.0, 0.0], [1.0, 0.0]]] * num_states, dtype=tf.float32), reinterpreted_batch_ndims=1),
        tfd.Blockwise([
            tfd.Categorical(probs=[[0.0, 1.0]] * num_states),
            tfd.VonMises(loc=vm_locs, concentration=vm_cons),
            tfd.Gamma(concentration=gamma_shapes, rate=gamma_rates)
        ], dtype_override=tf.float32)
    ]
)

hmm3 = tfd.HiddenMarkovModel(
    initial_distribution = tfd.Categorical(probs=initial_probs),
    transition_distribution = tfd.Categorical(probs=transition_probs),
    observation_distribution = mixed_dists,
    num_steps = 15
)

hmm3.log_prob(hmm3.sample()) # error here

The error I am getting is this:
---------------------------------------------------------------------------
InvalidArgumentError                      Traceback (most recent call last)
<ipython-input-638-bade53a804d5> in <module>
----> 1 hmm3.log_prob(observations[:15])
~\anaconda3\envs\shark_research_project\lib\site-packages\tensorflow_probability\python\distributions\distribution.py in log_prob(self, value, name, **kwargs)
   1294         values of type `self.dtype`.
   1295     """
-> 1296     return self._call_log_prob(value, name, **kwargs)
   1297 
   1298   def _call_prob(self, value, name, **kwargs):

~\anaconda3\envs\shark_research_project\lib\site-packages\tensorflow_probability\python\distributions\distribution.py in _call_log_prob(self, value, name, **kwargs)
   1276     with self._name_and_control_scope(name, value, kwargs):
   1277       if hasattr(self, '_log_prob'):
-> 1278         return self._log_prob(value, **kwargs)
   1279       if hasattr(self, '_prob'):
   1280         return tf.math.log(self._prob(value, **kwargs))

~\anaconda3\envs\shark_research_project\lib\site-packages\tensorflow_probability\python\distributions\hidden_markov_model.py in _log_prob(self, value)
    487       # working_obs :: num_steps batch_shape 1 underlying_event_shape
    488 
--> 489       observation_probs = observation_distribution.log_prob(working_obs)
    490       # observation_probs :: num_steps batch_shape num_states
    491 

~\anaconda3\envs\shark_research_project\lib\site-packages\tensorflow_probability\python\distributions\distribution.py in log_prob(self, value, name, **kwargs)
   1294         values of type `self.dtype`.
   1295     """
-> 1296     return self._call_log_prob(value, name, **kwargs)
   1297 
   1298   def _call_prob(self, value, name, **kwargs):

~\anaconda3\envs\shark_research_project\lib\site-packages\tensorflow_probability\python\distributions\distribution.py in _call_log_prob(self, value, name, **kwargs)
   1276     with self._name_and_control_scope(name, value, kwargs):
   1277       if hasattr(self, '_log_prob'):
-> 1278         return self._log_prob(value, **kwargs)
   1279       if hasattr(self, '_prob'):
   1280         return tf.math.log(self._prob(value, **kwargs))

~\anaconda3\envs\shark_research_project\lib\site-packages\tensorflow_probability\python\distributions\mixture.py in _log_prob(self, x)
    281   def _log_prob(self, x):
    282     x = tf.convert_to_tensor(x, name='x')
--> 283     distribution_log_probs = [d.log_prob(x) for d in self.components]
    284     cat_log_probs = self._cat_probs(log_probs=True)
    285     final_log_probs = [

~\anaconda3\envs\shark_research_project\lib\site-packages\tensorflow_probability\python\distributions\mixture.py in <listcomp>(.0)
    281   def _log_prob(self, x):
    282     x = tf.convert_to_tensor(x, name='x')
--> 283     distribution_log_probs = [d.log_prob(x) for d in self.components]
    284     cat_log_probs = self._cat_probs(log_probs=True)
    285     final_log_probs = [

~\anaconda3\envs\shark_research_project\lib\site-packages\tensorflow_probability\python\distributions\distribution.py in log_prob(self, value, name, **kwargs)
   1294         values of type `self.dtype`.
   1295     """
-> 1296     return self._call_log_prob(value, name, **kwargs)
   1297 
   1298   def _call_prob(self, value, name, **kwargs):

~\anaconda3\envs\shark_research_project\lib\site-packages\tensorflow_probability\python\distributions\distribution.py in _call_log_prob(self, value, name, **kwargs)
   1276     with self._name_and_control_scope(name, value, kwargs):
   1277       if hasattr(self, '_log_prob'):
-> 1278         return self._log_prob(value, **kwargs)
   1279       if hasattr(self, '_prob'):
   1280         return tf.math.log(self._prob(value, **kwargs))

~\anaconda3\envs\shark_research_project\lib\site-packages\tensorflow_probability\python\distributions\independent.py in _log_prob(self, x, **kwargs)
    283   def _log_prob(self, x, **kwargs):
    284     return self._reduce(
--> 285         self._sum_fn(), self.distribution.log_prob(x, **kwargs))
    286 
    287   def _unnormalized_log_prob(self, x, **kwargs):

~\anaconda3\envs\shark_research_project\lib\site-packages\tensorflow_probability\python\distributions\distribution.py in log_prob(self, value, name, **kwargs)
   1294         values of type `self.dtype`.
   1295     """
-> 1296     return self._call_log_prob(value, name, **kwargs)
   1297 
   1298   def _call_prob(self, value, name, **kwargs):

~\anaconda3\envs\shark_research_project\lib\site-packages\tensorflow_probability\python\distributions\distribution.py in _call_log_prob(self, value, name, **kwargs)
   1276     with self._name_and_control_scope(name, value, kwargs):
   1277       if hasattr(self, '_log_prob'):
-> 1278         return self._log_prob(value, **kwargs)
   1279       if hasattr(self, '_prob'):
   1280         return tf.math.log(self._prob(value, **kwargs))

~\anaconda3\envs\shark_research_project\lib\site-packages\tensorflow_probability\python\distributions\categorical.py in _log_prob(self, k)
    293     k, logits = _broadcast_cat_event_and_params(
    294         k, logits, base_dtype=dtype_util.base_dtype(self.dtype))
--> 295     return -tf.nn.sparse_softmax_cross_entropy_with_logits(
    296         labels=k, logits=logits)
    297 

~\anaconda3\envs\shark_research_project\lib\site-packages\tensorflow\python\util\dispatch.py in wrapper(*args, **kwargs)
    204     """Call target, and fall back on dispatchers if there is a TypeError."""
    205     try:
--> 206       return target(*args, **kwargs)
    207     except (TypeError, ValueError):
    208       # Note: convert_to_eager_tensor currently raises a ValueError, not a

~\anaconda3\envs\shark_research_project\lib\site-packages\tensorflow\python\ops\nn_ops.py in sparse_softmax_cross_entropy_with_logits_v2(labels, logits, name)
   4226       of the labels is not equal to the rank of the logits minus one.
   4227   """
-> 4228   return sparse_softmax_cross_entropy_with_logits(
   4229       labels=labels, logits=logits, name=name)
   4230 

~\anaconda3\envs\shark_research_project\lib\site-packages\tensorflow\python\util\dispatch.py in wrapper(*args, **kwargs)
    204     """Call target, and fall back on dispatchers if there is a TypeError."""
    205     try:
--> 206       return target(*args, **kwargs)
    207     except (TypeError, ValueError):
    208       # Note: convert_to_eager_tensor currently raises a ValueError, not a

~\anaconda3\envs\shark_research_project\lib\site-packages\tensorflow\python\ops\nn_ops.py in sparse_softmax_cross_entropy_with_logits(_sentinel, labels, logits, name)
   4159       # The second output tensor contains the gradients.  We use it in
   4160       # _CrossEntropyGrad() in nn_grad but not here.
-> 4161       cost, _ = gen_nn_ops.sparse_softmax_cross_entropy_with_logits(
   4162           precise_logits, labels, name=name)
   4163       cost = array_ops.reshape(cost, labels_shape)

~\anaconda3\envs\shark_research_project\lib\site-packages\tensorflow\python\ops\gen_nn_ops.py in sparse_softmax_cross_entropy_with_logits(features, labels, name)
  11247       return _result
  11248     except _core._NotOkStatusException as e:
> 11249       _ops.raise_from_not_ok_status(e, name)
  11250     except _core._FallbackException:
  11251       pass

~\anaconda3\envs\shark_research_project\lib\site-packages\tensorflow\python\framework\ops.py in raise_from_not_ok_status(e, name)
   6895   message = e.message + (" name: " + name if name is not None else "")
   6896   # pylint: disable=protected-access
-> 6897   six.raise_from(core._status_to_exception(e.code, message), None)
   6898   # pylint: enable=protected-access
   6899 

~\anaconda3\envs\shark_research_project\lib\site-packages\six.py in raise_from(value, from_value)

InvalidArgumentError: Received a label value of -1 which is outside the valid range of [0, 2).  Label values: 1 0 0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 1 0 0 1 0 0 1 0 0 1 0 0 0 0 0 0 0 0 1 0 0 1 0 0 1 0 0 1 0 0 1 1 0 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 1 1 0 1 1 0 1 0 0 1 0 0 0 0 0 0 0 0 1 0 -1 1 0 -1 [Op:SparseSoftmaxCrossEntropyWithLogits]

The strangest part is that when I create the hidden Markov model with 14 steps and call hmm3.log_prob(hmm3.sample), I don’t get any errors. Please let me know if I should clarify anything or include more information. Any help would be appreciated!