Thanks, this solved the issue I was running into! Once again I’ve learned something new after having bashed my head against a wall for a while trying to use the wrong tool! Unfortunately, I think this solution has brought forward another issue. When I try to fit the HMM using the Adam optimizer to maximize the log probability, I get nan gradients for the gamma concentrations after the first iteration. You mentioned that the input to log_prob is evaluated under all mixture components, which leads me to believe that the issue is with the gamma concentrations being updated to 1.1 on the first training step, after which observations of the form [0, 0, 0] have a log probability of -Inf under the gamma distribution. Then again, I could be totally off, I am not very familiar with Tensorflow. Here is the relevant code:
# Define a loss function
def log_prob():
    return hmm3.log_prob(observations)

# Define an optimizer to perform back propagation
optimizer = tf.keras.optimizers.Adam(learning_rate=0.1)

# Make sure probabilities sum to 1
def normalize_probs(probs):
    abs_probs = tf.math.abs(probs)
    if len(probs.shape) > 1:
        sums = tf.reshape(tf.reduce_sum(abs_probs, axis=1), [probs.shape[0], 1])
    else:
        sums = tf.reduce_sum(abs_probs)
    return abs_probs / sums

def wrap_to_pi(A):
    return ((A - np.pi) % (2 * np.pi) - np.pi)

# Run a step of the optimizer
@tf.function(autograph=False)
def train_op():
    with tf.GradientTape() as tape:
        neg_log_prob = -log_prob()
    vars = [initial_probs, transition_probs, vm_locs, vm_cons, gamma_shapes, gamma_rates]
    grads = tape.gradient(neg_log_prob, vars)
    optimizer.apply_gradients(zip(grads, vars))
    initial_probs.assign(normalize_probs(initial_probs))
    transition_probs.assign(normalize_probs(transition_probs))
    vm_locs.assign(wrap_to_pi(vm_locs))
    vm_cons.assign(tf.math.abs(vm_cons))
    gamma_shapes.assign(tf.math.abs(gamma_shapes))
    gamma_rates.assign(tf.math.abs(gamma_rates))
    return (neg_log_prob, *vars), grads

# Train on the observations
loss_history = []
for step in range(10):
    ts, grads = train_op()
    loss, ip, tp, vl, vc, gs, gr = [t.numpy() for t in ts]
    loss_history.append(loss)
    if step % 1 == 0:
        print("step {}: log prob {}\nInitial probs: {}\nTransition probs:\n{}\nVon Mises locs: {}\nVon Mises cons: {}\nGamma shapes: {}\nGamma rates: {}\n".format(step, -loss, ip, tp, vl, vc, gs, gr))

I would appreciate any insight you might have on how to fix this new error, and thanks a lot for the help you’ve provided so far!