Fix test cases for k2.union() (#853)

Author: csukuangfj

k2/python/k2/__init__.py

@@ -21,7 +21,6 @@
 #
 from .autograd import intersect_dense
 from .autograd import intersect_dense_pruned
-from .autograd import union
 from .ctc_loss import CtcLoss
 from .ctc_loss import ctc_loss
 from .dense_fsa_vec import DenseFsaVec
@@ -51,6 +50,7 @@
 from .fsa_algo import replace_fsa
 from .fsa_algo import shortest_path
 from .fsa_algo import top_sort
+from .fsa_algo import union
 from .fsa_properties import to_str as properties_to_str
 from .nbest import Nbest
 from .ops import cat

k2/python/k2/autograd.py

@@ -645,53 +645,6 @@ def backward(ctx, out_fsa_grad: torch.Tensor) \
         )
 
 
-class _UnionFunction(torch.autograd.Function):
-
-    @staticmethod
-    def forward(ctx, fsas: Fsa, out_fsa: List[Fsa],
-                unused_fsas_scores: torch.Tensor) -> torch.Tensor:
-        '''Compute the union of all fsas in a FsaVec.
-
-        Args:
-          fsas:
-            The input FsaVec. Caution: We require that each fsa in the FsaVec
-            is non-empty (i.e., with at least two states).
-          out_fsa:
-            A list containing one entry. Since this function can only return
-            values of type `torch.Tensor`, we return the union result in the
-            list.
-          unused_fsas_scores:
-            It is the same as `fsas.scores`, whose sole purpose is for autograd.
-            It is not used in this function.
-        '''
-        need_arc_map = True
-        ragged_arc, arc_map = _k2.union(fsas.arcs, need_arc_map)
-        out_fsa[0] = Fsa(ragged_arc)
-
-        for name, value in fsas.named_tensor_attr(include_scores=False):
-            value = k2.index(value, arc_map)
-            setattr(out_fsa[0], name, value)
-
-        for name, value in fsas.named_non_tensor_attr():
-            setattr(out_fsa[0], name, value)
-        ctx.arc_map = arc_map
-        ctx.save_for_backward(unused_fsas_scores)
-
-        return out_fsa[0].scores  # the return value will be discarded
-
-    @staticmethod
-    def backward(ctx, out_fsa_grad: torch.Tensor
-                ) -> Tuple[None, None, torch.Tensor]:  # noqa
-        arc_map = ctx.arc_map
-        fsas_scores, = ctx.saved_tensors
-        ans = torch.zeros(fsas_scores.size(0),
-                          dtype=torch.float32,
-                          device=fsas_scores.device,
-                          requires_grad=False)
-        _k2.index_add(arc_map, out_fsa_grad, ans)
-        return None, None, ans
-
-
 def intersect_dense_pruned(a_fsas: Fsa,
                            b_fsas: DenseFsaVec,
                            search_beam: float,
@@ -843,23 +796,3 @@ def intersect_dense(a_fsas: Fsa,
                                   a_fsas.scores, b_fsas.scores, a_to_b_map,
                                   seqframe_idx_name, frame_idx_name)
     return out_fsa[0]
-
-
-def union(fsas: Fsa) -> Fsa:
-    '''Compute the union of a FsaVec.
-
-    Caution:
-      We require that every fsa in fsas is non-empty, i.e.,
-      contains at least two states
-
-    Args:
-      fsas:
-        A FsaVec. That is, len(fsas.shape) == 3.
-
-    Returns:
-      A single Fsa that is the union of the input fsas.
-    '''
-
-    out_fsa = [0]  # as a placeholder
-    _UnionFunction.apply(fsas, out_fsa, fsas.scores)
-    return out_fsa[0]

k2/python/k2/fsa_algo.py

@@ -1179,3 +1179,24 @@ def levenshtein_alignment(
         alignment, "__ins_del_score_offset_internal_attr_")
 
     return alignment
+
+
+def union(fsas: Fsa) -> Fsa:
+    '''Compute the union of a FsaVec.
+
+    Caution:
+      We require that every fsa in fsas is non-empty, i.e.,
+      contains at least two states
+
+    Args:
+      fsas:
+        A FsaVec. That is, len(fsas.shape) == 3.
+
+    Returns:
+      A single Fsa that is the union of the input fsas.
+    '''
+    need_arc_map = True
+    ragged_arc, arc_map = _k2.union(fsas.arcs, need_arc_map)
+
+    out_fsa = k2.utils.fsa_from_unary_function_tensor(fsas, ragged_arc, arc_map)
+    return out_fsa

k2/python/tests/union_test.py

@@ -63,9 +63,40 @@ def test(self):
             fsa1 = k2.Fsa.from_str(s1)
             fsa2 = k2.Fsa.from_str(s2)
 
+            fsa0.tensor_attr = torch.tensor([1, 2, 3, 4, 5, 6],
+                                            dtype=torch.int32,
+                                            device=device)
+            fsa0.ragged_tensor_attr = k2.RaggedTensor(
+                fsa0.tensor_attr.unsqueeze(-1))
+
+            fsa1.tensor_attr = torch.tensor([7],
+                                            dtype=torch.int32,
+                                            device=device)
+
+            fsa1.ragged_tensor_attr = k2.RaggedTensor(
+                fsa1.tensor_attr.unsqueeze(-1))
+
+            fsa2.tensor_attr = torch.tensor([8, 9, 10, 11],
+                                            dtype=torch.int32,
+                                            device=device)
+
+            fsa2.ragged_tensor_attr = k2.RaggedTensor(
+                fsa2.tensor_attr.unsqueeze(-1))
+
             fsa_vec = k2.create_fsa_vec([fsa0, fsa1, fsa2]).to(device)
 
             fsa = k2.union(fsa_vec)
+
+            expected_tensor_attr = torch.tensor(
+                [0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10,
+                 11]).to(fsa.tensor_attr)
+            assert torch.all(torch.eq(fsa.tensor_attr, expected_tensor_attr))
+
+            expected_ragged_tensor_attr = k2.RaggedTensor(
+                expected_tensor_attr.unsqueeze(-1)).remove_values_eq(0)
+            assert str(expected_ragged_tensor_attr) == str(
+                fsa.ragged_tensor_attr)
+
             assert torch.allclose(
                 fsa.arcs.values()[:, :3],
                 torch.tensor([