replace FSA with Array2 in RmEpsilons

Author: qindazhu

k2/csrc/CMakeLists.txt

@@ -8,6 +8,7 @@ add_library(fsa
   fsa_renderer.cc
   fsa_util.cc
   properties.cc
+  rmepsilon.cc
   util.cc
   weights.cc
 )
@@ -46,6 +47,7 @@ set(fsa_tests
   fsa_test
   fsa_util_test
   properties_test
+  rmepsilon_test
   weights_test
 )
 

k2/csrc/array.h

@@ -94,13 +94,13 @@ struct Array2 {
   using PtrT = Ptr;
   using ValueType = typename std::iterator_traits<Ptr>::value_type;
 
-  Array2() : size1(0), indexes(&size1), size2(0), data(nullptr) {}
-  Array2(IndexT size1, IndexT *indexes, IndexT size2, PtrT data)
-      : size1(size1), indexes(indexes), size2(size2), data(data) {}
-  void Init(IndexT size1, IndexT *indexes, IndexT size2, PtrT data) {
+  Array2() : size1(0), size2(0), indexes(&size1), data(nullptr) {}
+  Array2(IndexT size1, IndexT size2, IndexT *indexes, PtrT data)
+      : size1(size1), size2(size2), indexes(indexes), data(data) {}
+  void Init(IndexT size1, IndexT size2, IndexT *indexes, PtrT data) {
     this->size1 = size1;
-    this->indexes = indexes;
     this->size2 = size2;
+    this->indexes = indexes;
     this->data = data;
   }
 

k2/csrc/aux_labels.cc

@@ -245,7 +245,7 @@ void FstInverter::GetOutput(Fsa *fsa_out, AuxLabels *labels_out) {
 
   std::vector<int32_t> arc_map;
   ReorderArcs(arcs, fsa_out, &arc_map);
-  AuxLabels labels_tmp(labels_out->size1, start_pos.data(), labels_out->size2,
+  AuxLabels labels_tmp(labels_out->size1, labels_out->size2, start_pos.data(),
                        labels.data());
   AuxLabels1Mapper aux_mapper(labels_tmp, arc_map);
   // don't need to call `GetSizes` here as `labels_out` has been initialized

k2/csrc/aux_labels_test.cc

@@ -134,7 +134,7 @@ TEST(AuxLabels, InvertFst) {
     std::vector<int32_t> start_pos = {0, 1, 3, 6, 7};
     std::vector<int32_t> labels = {1, 2, 3, 4, 5, 6, 7};
     AuxLabels labels_in(static_cast<int32_t>(start_pos.size()) - 1,
-                        start_pos.data(), static_cast<int32_t>(labels.size()),
+                        static_cast<int32_t>(labels.size()), start_pos.data(),
                         labels.data());
 
     FstInverter fst_inverter(fsa_in, labels_in);
@@ -162,7 +162,7 @@ TEST(AuxLabels, InvertFst) {
     EXPECT_EQ(start_pos.size(), fsa_in.size2 + 1);
     std::vector<int32_t> labels = {1, 2, 3, 5, 6, 7, -1, -1, -1};
     AuxLabels labels_in(static_cast<int32_t>(start_pos.size()) - 1,
-                        start_pos.data(), static_cast<int32_t>(labels.size()),
+                        static_cast<int32_t>(labels.size()), start_pos.data(),
                         labels.data());
 
     FstInverter fst_inverter(fsa_in, labels_in);
@@ -212,7 +212,7 @@ TEST(AuxLabels, InvertFst) {
     EXPECT_EQ(start_pos.size(), fsa_in.size2 + 1);
     std::vector<int32_t> labels = {1, 2, 3, 5, 6, 7, 8, -1, 9, 10, -1};
     AuxLabels labels_in(static_cast<int32_t>(start_pos.size()) - 1,
-                        start_pos.data(), static_cast<int32_t>(labels.size()),
+                        static_cast<int32_t>(labels.size()), start_pos.data(),
                         labels.data());
 
     FstInverter fst_inverter(fsa_in, labels_in);

k2/csrc/determinize.h

@@ -172,8 +172,11 @@ struct MaxTracebackState {
                         // sequence of symbols we took to get here)
 
   // This constructor is for the start-state (state zero) of the input FSA.
-  MaxTracebackState()
-      : state_id(0), arc_id(-1), prev_state(nullptr), forward_prob(0.0) {}
+  explicit MaxTracebackState(int32_t state_id = 0, double forward_prob = 0.0)
+      : state_id(state_id),
+        arc_id(-1),
+        prev_state(nullptr),
+        forward_prob(forward_prob) {}
 
   /**
      @param [in] state_id  State in input FSA that this corresponds to

k2/csrc/fsa_algo.cc

@@ -52,78 +52,6 @@ inline int32_t InsertIntersectionState(
   return result.first->second;
 }
 
-/**
-   A TraceBack() function used in RmEpsilonsPrunedLogSum.  It finds derivative
-   information for all arcs in a sub-graph. Generally, in
-   RmEpsilonsPrunedLogSum, we actually get a sub-graph when we find a
-   non-epsilon arc starting from a particular state `s` (from which we are
-   trying to remove epsilon arcs). All leaving arcs of all states in this
-   sub-graph are epsilon arcs except the last one. Then, from the last state, we
-   need to trace back to state `s` to find the derivative information for all
-   epsilon arcs in this graph.
-       @param [in] curr_states   (This is consumed destructively, i.e. don't
-                       expect it to contain the same set on exit).
-                       A set of states, stored as a std::map that mapping
-                       state_id in input FSA to the corresponding
-                       LogSumTracebackState we created for this state;
-                       we'll iteratively trace back this set one element
-                       (processing all entering arcs) at a time.  At entry
-                       it must have size() == 1 which contains the last
-                       state mentioned above; it will also have size() == 1
-                       at exit which contains the state `s` above.
-       @param [in] arc_weights_in  Weights on the arcs of the input FSA
-       @param [out] deriv_out  Some derivative information at the output
-                       will be written to here, which tells us how the weight
-                       of the non-epsilon arc we created from the above
-                       sub-graph varies as a function of the weights on the
-                       arcs of the input FSA; it's a list
-                       (input_arc_id, deriv) where, mathematically,
-                       0 < deriv <= 1 (but we might still get exact zeros
-                       due to limitations of floating point representation).
- */
-static void TraceBackRmEpsilonsLogSum(
-    std::map<int32_t, k2::LogSumTracebackState *> *curr_states,
-    const float *arc_weights_in,
-    std::vector<std::pair<int32_t, float>> *deriv_out) {
-  CHECK_EQ(curr_states->size(), 1);
-  deriv_out->clear();
-
-  // as the input fsa is top-sorted, we traverse states in a reverse order so we
-  // can process them when they already have correct backward_prob (all leaving
-  // arcs have been processed).
-  k2::LogSumTracebackState *state_ptr = curr_states->rbegin()->second;
-  // In the standard forward-backward algorithm for HMMs this backward_prob
-  // would, mathematically, be 0.0, but if we set it to the negative of the
-  // forward prob we can avoid having to subtract the total log-prob
-  // when we compute posterior/occupation probabilities for arcs.
-  state_ptr->backward_prob = -state_ptr->forward_prob;
-  while (!state_ptr->prev_elements.empty()) {
-    double backward_prob = state_ptr->backward_prob;
-    for (const auto &link : state_ptr->prev_elements) {
-      auto arc_log_posterior =
-          static_cast<float>(link.forward_prob + backward_prob);
-      deriv_out->emplace_back(link.arc_index, expf(arc_log_posterior));
-      k2::LogSumTracebackState *prev_state = link.prev_state.get();
-      double new_backward_prob = backward_prob + arc_weights_in[link.arc_index];
-      auto result = curr_states->emplace(prev_state->state_id, prev_state);
-      if (result.second) {
-        prev_state->backward_prob = new_backward_prob;
-      } else {
-        prev_state->backward_prob =
-            k2::LogAdd(new_backward_prob, prev_state->backward_prob);
-      }
-    }
-    // we have processed all entering arcs of state curr_states->rbegin(),
-    // we'll remove it now. As std::map.erase() does not support passing a
-    // reverse iterator, we here pass --end();
-    curr_states->erase(--curr_states->end());
-    CHECK(!curr_states->empty());
-    state_ptr = curr_states->rbegin()->second;
-  }
-  // we have reached the state from which we are trying to remove epsilon arcs.
-  CHECK_EQ(curr_states->size(), 1);
-}
-
 }  // namespace
 
 namespace k2 {
@@ -350,215 +278,6 @@ bool Connect(const Fsa &a, Fsa *b, std::vector<int32_t> *arc_map /*=nullptr*/) {
   return is_acyclic;
 }
 
-void RmEpsilonsPrunedMax(const WfsaWithFbWeights &a, float beam, Fsa *b,
-                         std::vector<std::vector<int32_t>> *arc_derivs) {
-  CHECK_EQ(a.weight_type, kMaxWeight);
-  CHECK_GT(beam, 0);
-  CHECK_NOTNULL(b);
-  CHECK_NOTNULL(arc_derivs);
-  b->arc_indexes.clear();
-  b->arcs.clear();
-  arc_derivs->clear();
-
-  const auto &fsa = a.fsa;
-  if (IsEmpty(fsa)) return;
-  int32_t num_states_a = fsa.NumStates();
-  int32_t final_state = fsa.FinalState();
-  const auto &arcs_a = fsa.data;
-  const float *arc_weights_a = a.arc_weights;
-
-  // identify all states that should be kept
-  std::vector<char> non_eps_in(num_states_a, 0);
-  non_eps_in[0] = 1;
-  for (const auto &arc : fsa) {
-    // We suppose the input fsa `a` is top-sorted, but only check this in DEBUG
-    // time.
-    DCHECK_GE(arc.dest_state, arc.src_state);
-    if (arc.label != kEpsilon) non_eps_in[arc.dest_state] = 1;
-  }
-
-  // remap state id
-  std::vector<int32_t> state_map_a2b(num_states_a, -1);
-  int32_t num_states_b = 0;
-  for (int32_t i = 0; i != num_states_a; ++i) {
-    if (non_eps_in[i] == 1) state_map_a2b[i] = num_states_b++;
-  }
-  b->arc_indexes.reserve(num_states_b + 1);
-  int32_t arc_num_b = 0;
-
-  const double *forward_state_weights = a.ForwardStateWeights();
-  const double *backward_state_weights = a.BackwardStateWeights();
-  const double best_weight = forward_state_weights[final_state] - beam;
-  for (int32_t i = 0; i != num_states_a; ++i) {
-    if (non_eps_in[i] != 1) continue;
-    b->arc_indexes.push_back(arc_num_b);
-    int32_t curr_state_b = state_map_a2b[i];
-    // as the input FSA is top-sorted, we use a map here so we can process
-    // states when they already have the best cost they are going to get
-    std::map<int32_t, double>
-        local_forward_weights;  // state -> local_forward_state_weights of this
-                                // state
-    // state -> (src_state, arc_index) entering this state which contributes to
-    // `local_forward_weights` of this state.
-    std::unordered_map<int32_t, std::pair<int32_t, int32_t>>
-        local_backward_arcs;
-    local_forward_weights.emplace(i, forward_state_weights[i]);
-    // `-1` means we have traced back to current state `i`
-    local_backward_arcs.emplace(i, std::make_pair(i, -1));
-    while (!local_forward_weights.empty()) {
-      std::pair<int32_t, double> curr_local_forward_weights =
-          *(local_forward_weights.begin());
-      local_forward_weights.erase(local_forward_weights.begin());
-      int32_t state = curr_local_forward_weights.first;
-
-      int32_t arc_end = fsa.indexes[state + 1];
-      for (int32_t arc_index = fsa.indexes[state]; arc_index != arc_end;
-           ++arc_index) {
-        int32_t next_state = arcs_a[arc_index].dest_state;
-        int32_t label = arcs_a[arc_index].label;
-        double next_weight =
-            curr_local_forward_weights.second + arc_weights_a[arc_index];
-        if (next_weight + backward_state_weights[next_state] >= best_weight) {
-          if (label == kEpsilon) {
-            auto result =
-                local_forward_weights.emplace(next_state, next_weight);
-            if (result.second) {
-              local_backward_arcs[next_state] =
-                  std::make_pair(state, arc_index);
-            } else {
-              if (next_weight > result.first->second) {
-                result.first->second = next_weight;
-                local_backward_arcs[next_state] =
-                    std::make_pair(state, arc_index);
-              }
-            }
-          } else {
-            b->arcs.emplace_back(curr_state_b, state_map_a2b[next_state],
-                                 label);
-            std::vector<int32_t> curr_arc_deriv;
-            std::pair<int32_t, int32_t> curr_backward_arc{state, arc_index};
-            auto *backward_arc = &curr_backward_arc;
-            while (backward_arc->second != -1) {
-              curr_arc_deriv.push_back(backward_arc->second);
-              backward_arc = &(local_backward_arcs[backward_arc->first]);
-            }
-            std::reverse(curr_arc_deriv.begin(), curr_arc_deriv.end());
-            arc_derivs->emplace_back(std::move(curr_arc_deriv));
-            ++arc_num_b;
-          }
-        }
-      }
-    }
-  }
-  // duplicate of final state
-  b->arc_indexes.push_back(b->arc_indexes.back());
-}
-
-void RmEpsilonsPrunedLogSum(
-    const WfsaWithFbWeights &a, float beam, Fsa *b,
-    std::vector<float> *b_arc_weights,
-    std::vector<std::vector<std::pair<int32_t, float>>> *arc_derivs) {
-  CHECK_GT(beam, 0);
-  CHECK_NOTNULL(b);
-  CHECK_NOTNULL(b_arc_weights);
-  CHECK_NOTNULL(arc_derivs);
-  b->arc_indexes.clear();
-  b->arcs.clear();
-  b_arc_weights->clear();
-  arc_derivs->clear();
-
-  const auto &fsa = a.fsa;
-  if (IsEmpty(fsa)) return;
-  int32_t num_states_a = fsa.NumStates();
-  int32_t final_state = fsa.FinalState();
-  const auto &arcs_a = fsa.data;
-  const float *arc_weights_a = a.arc_weights;
-
-  // identify all states that should be kept
-  std::vector<char> non_eps_in(num_states_a, 0);
-  non_eps_in[0] = 1;
-  for (const auto &arc : fsa) {
-    // We suppose the input fsa `a` is top-sorted, but only check this in DEBUG
-    // time.
-    DCHECK_GE(arc.dest_state, arc.src_state);
-    if (arc.label != kEpsilon) non_eps_in[arc.dest_state] = 1;
-  }
-
-  // remap state id
-  std::vector<int32_t> state_map_a2b(num_states_a, -1);
-  int32_t num_states_b = 0;
-  for (int32_t i = 0; i != num_states_a; ++i) {
-    if (non_eps_in[i] == 1) state_map_a2b[i] = num_states_b++;
-  }
-  b->arc_indexes.reserve(num_states_b + 1);
-  int32_t arc_num_b = 0;
-
-  const double *forward_state_weights = a.ForwardStateWeights();
-  const double *backward_state_weights = a.BackwardStateWeights();
-  const double best_weight = forward_state_weights[final_state] - beam;
-  for (int32_t i = 0; i != num_states_a; ++i) {
-    if (non_eps_in[i] != 1) continue;
-    b->arc_indexes.push_back(arc_num_b);
-    int32_t curr_state_b = state_map_a2b[i];
-    // as the input FSA is top-sorted, we use a set here so we can process
-    // states when they already have costs over all paths they are going to get
-    std::set<int32_t> qstates;
-    std::unordered_map<int32_t, std::shared_ptr<LogSumTracebackState>>
-        traceback_states;  // state -> LogSumTracebackState of this state
-    std::shared_ptr<LogSumTracebackState> start_state(
-        new LogSumTracebackState(i, forward_state_weights[i]));
-    double start_forward_weights = start_state->forward_prob;
-    traceback_states.emplace(i, start_state);
-    qstates.insert(i);
-    while (!qstates.empty()) {
-      int32_t state = *(qstates.begin());
-      qstates.erase(qstates.begin());
-
-      const auto &curr_traceback_state = traceback_states[state];
-      double curr_forward_weights = curr_traceback_state->forward_prob;
-      int32_t arc_end = fsa.indexes[state + 1];
-      for (int32_t arc_index = fsa.indexes[state]; arc_index != arc_end;
-           ++arc_index) {
-        int32_t next_state = arcs_a[arc_index].dest_state;
-        int32_t label = arcs_a[arc_index].label;
-        float curr_arc_weight = arc_weights_a[arc_index];
-        double next_weight = curr_forward_weights + curr_arc_weight;
-        if (next_weight + backward_state_weights[next_state] >= best_weight) {
-          if (label == kEpsilon) {
-            auto result = traceback_states.emplace(next_state, nullptr);
-            if (result.second) {
-              result.first->second = std::make_shared<LogSumTracebackState>(
-                  next_state, curr_traceback_state, arc_index, curr_arc_weight);
-              qstates.insert(next_state);
-            } else {
-              result.first->second->Accept(curr_traceback_state, arc_index,
-                                           curr_arc_weight);
-            }
-          } else {
-            b->arcs.emplace_back(curr_state_b, state_map_a2b[next_state],
-                                 label);
-            b_arc_weights->push_back(curr_forward_weights + curr_arc_weight -
-                                     start_forward_weights);
-
-            std::vector<std::pair<int32_t, float>> curr_arc_deriv;
-            std::map<int32_t, LogSumTracebackState *> curr_states;
-            curr_states.emplace(state, curr_traceback_state.get());
-            TraceBackRmEpsilonsLogSum(&curr_states, arc_weights_a,
-                                      &curr_arc_deriv);
-            std::reverse(curr_arc_deriv.begin(), curr_arc_deriv.end());
-            // push derivs info of current arc
-            curr_arc_deriv.emplace_back(arc_index, 1);
-            arc_derivs->emplace_back(std::move(curr_arc_deriv));
-            ++arc_num_b;
-          }
-        }
-      }
-    }
-  }
-  // duplicate of final state
-  b->arc_indexes.push_back(b->arc_indexes.back());
-}
-
 bool Intersect(const Fsa &a, const Fsa &b, Fsa *c,
                std::vector<int32_t> *arc_map_a /*= nullptr*/,
                std::vector<int32_t> *arc_map_b /*= nullptr*/) {