<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 def predict(self, parameters, viterbi):
 """ Run forward algorithm to find the predicted distribution over classes. """

        x_dot_parameters = np.einsum('ijk,kl->ijl', self.x, parameters)

        if not viterbi:
            alpha = forward_predict(self._lattice, x_dot_parameters,
                                    self.state_machine.n_states)
        else:
            alpha = forward_max_predict(self._lattice, x_dot_parameters,
                                        self.state_machine.n_states)

        I, J, _ = self.x.shape

        class_Z = {}
        Z = -np.inf

        for state, predicted_class in self.states_to_classes.items():
            weight = alpha[I - 1, J - 1, state]
            class_Z[self.states_to_classes[state]] = weight
            Z = np.logaddexp(Z, weight)

        return {label: np.exp(class_z - Z) for label, class_z in class_Z.items()}