def test_build_event(self):
        evt = event.build_event("evt_name", arg1="val1", arg2=None)
        assert evt.name == "evt_name"
        assert evt.args == {'arg1': "val1", 'arg2': None}
def update_state(self, y_true, y_pred, sample_weight=None):
    sample_weight = _apply_mask(y_true, sample_weight, self._masked_tokens,
                                self._dtype)
    num_classes = tf.shape(y_pred)[-1]
    y_true = tf.reshape(y_true, [-1])
    y_pred = tf.reshape(y_pred, [-1, num_classes])
    sample_weight = tf.reshape(sample_weight, [-1])
    super().update_state(y_true, y_pred, sample_weight)