add beam in RandEquivalent

Author: qindazhu

k2/csrc/fsa_equivalent.cc

@@ -169,8 +169,11 @@ bool IsRandEquivalent(const Fsa &a, const Fsa &b, std::size_t npath /*=100*/) {
 
 template <FbWeightType Type>
 bool IsRandEquivalent(const Fsa &a, const float *a_weights, const Fsa &b,
-                      const float *b_weights, bool top_sorted /*=true*/,
-                      std::size_t npath /*= 100*/) {
+                      const float *b_weights, float beam /*=kFloatInfinity*/,
+                      bool top_sorted /*=true*/, std::size_t npath /*= 100*/) {
+  CHECK_GT(beam, 0);
+  CHECK_NOTNULL(a_weights);
+  CHECK_NOTNULL(b_weights);
   Fsa connected_a, connected_b, valid_a, valid_b;
   std::vector<int32_t> connected_a_arc_map, connected_b_arc_map,
       valid_a_arc_map, valid_b_arc_map;
@@ -199,6 +202,13 @@ bool IsRandEquivalent(const Fsa &a, const float *a_weights, const Fsa &b,
        (*(labels_difference.begin())) != kEpsilon))
     return false;
 
+  double loglike_cutoff;
+  if (beam != kFloatInfinity)
+    loglike_cutoff =
+        ShortestDistance<Type>(valid_a, valid_a_weights.data()) - beam;
+  else
+    loglike_cutoff = kDoubleNegativeInfinity;
+
   std::random_device rd;
   std::mt19937 gen(rd());
   std::bernoulli_distribution coin(0.5);
@@ -222,6 +232,7 @@ bool IsRandEquivalent(const Fsa &a, const float *a_weights, const Fsa &b,
     CHECK(top_sorted);
     double sum_a =
         ShortestDistance<Type>(a_compose_path, a_compose_weights.data());
+    if (sum_a < loglike_cutoff) sum_a = kDoubleNegativeInfinity;
     double sum_b =
         ShortestDistance<Type>(b_compose_path, b_compose_weights.data());
     if (!DoubleApproxEqual(sum_a, sum_b)) return false;
@@ -232,10 +243,11 @@ bool IsRandEquivalent(const Fsa &a, const float *a_weights, const Fsa &b,
 // explicit instantiation here
 template bool IsRandEquivalent<kMaxWeight>(const Fsa &a, const float *a_weights,
                                            const Fsa &b, const float *b_weights,
-                                           bool top_sorted, std::size_t npath);
+                                           float beam, bool top_sorted,
+                                           std::size_t npath);
 template bool IsRandEquivalent<kLogSumWeight>(
     const Fsa &a, const float *a_weights, const Fsa &b, const float *b_weights,
-    bool top_sorted, std::size_t npath);
+    float beam, bool top_sorted, std::size_t npath);
 
 bool RandomPath(const Fsa &a, Fsa *b,
                 std::vector<int32_t> *state_map /*=nullptr*/) {

k2/csrc/fsa_equivalent.h

@@ -32,16 +32,20 @@ bool IsRandEquivalent(const Fsa &a, const Fsa &b, std::size_t npath = 100);
   @param [in]  a_weights  Arc weights of `a`
   @param [in]  b          The other FSA to be checked the equivalence
   @param [in]  b_weights  Arc weights of `b`
-  @param [in]  top_sorted If both `a` and `b` are topological sorted or not.
-                          We may remove this flag if we finally find out that
-                          input FSAs in all scenarios are top-sorted.
+  @param [in]  beam       beam > 0 that affects pruning in `a`; the algorithm
+                          will keep paths that are within `beam` of the best
+                          path in `a` to check the equivalence. Just keep
+                          `kFloatInfinity` if you don't want pruning.
+  @param [in]  top_sorted The user may set this to true if both `a` and `b` are
+                          topologically sorted; this makes this function faster.
+                          Otherwise it must be set to false.
   @param [in]  npath      The number of paths will be generated to check the
                           equivalence of `a` and `b`
  */
 template <FbWeightType Type>
 bool IsRandEquivalent(const Fsa &a, const float *a_weights, const Fsa &b,
-                      const float *b_weights, bool top_sorted = true,
-                      std::size_t npath = 100);
+                      const float *b_weights, float beam = kFloatInfinity,
+                      bool top_sorted = true, std::size_t npath = 100);
 
 /*
   Gets a random path from an Fsa `a`, returns true if we get one path
@@ -69,10 +73,10 @@ bool RandomPathWithoutEpsilonArc(const Fsa &a, Fsa *b,
   @param [in] a_weights Arc weights of `a`
   @param [in] b    The other FSA to be intersected  Must satisfy
                    ArcSorted(b) and IsEpsilonFree(b). It is usually a path
-  generated from `RandomNonEpsilonPath`
+                   generated from `RandomNonEpsilonPath`
   @param [out] c   The composed FSA will be output to here.
   @param [out] c_weights Arc weights of output FSA `c` which are corresponding
-  arc weights in `a`
+                         arc weights in `a`
   @param [out] arc_map_a   If non-NULL, at exit will be a vector of
                    size c->arcs.size(), saying for each arc in
                    `c` what the source arc in `a` was, `-1` represents

k2/csrc/fsa_equivalent_test.cc

@@ -149,6 +149,13 @@ TEST(FsaEquivalent, IsWfsaRandEquivalent) {
                                                   c_weights.data());
     EXPECT_FALSE(status);
   }
+
+  // check equivalence with beam
+  {
+    bool status = IsRandEquivalent<kMaxWeight>(a, a_weights.data(), b,
+                                               b_weights.data(), 3);
+    EXPECT_FALSE(status);
+  }
 }
 
 TEST(FsaEquivalent, RandomPathFail) {

k2/csrc/weights.h

@@ -19,8 +19,10 @@
 
 namespace k2 {
 
+constexpr float kFloatInfinity = std::numeric_limits<float>::infinity();
 constexpr float kFloatNegativeInfinity =
     -std::numeric_limits<float>::infinity();
+constexpr double kDoubleInfinity = std::numeric_limits<double>::infinity();
 constexpr double kDoubleNegativeInfinity =
     -std::numeric_limits<double>::infinity();