def test_remove_axis_ragged_array(self): s = ''' [ [ [ 1 2 ] [ 0 ] ] [ [3 0 ] [ 2 ] ] ] ''' src = k2.RaggedInt(s) ans = k2.ragged.remove_axis(src, 0) self.assertEqual(str(ans), '[ [ 1 2 ] [ 0 ] [ 3 0 ] [ 2 ] ]') ans = k2.ragged.remove_axis(src, 1) self.assertEqual(str(ans), '[ [ 1 2 0 ] [ 3 0 2 ] ]')
def test_unique_sequences_two_axes(self): for device in self.devices: ragged = k2.RaggedInt( '[[1 3] [1 2] [1 2] [1 4] [1 3] [1 2] [1]]').to(device) unique, num_repeats, new2old = k2.ragged.unique_sequences( ragged, need_num_repeats=True, need_new2old_indexes=True) # [1, 3] has a larger hash value than [1, 2]; after sorting, # [1, 3] is placed after [1, 2] expected = k2.RaggedInt('[[1] [1 2] [1 3] [1 4]]') self.assertEqual(str(unique), str(expected)) expected_num_repeats = k2.RaggedInt('[[1 3 2 1]]') self.assertEqual(str(num_repeats), str(expected_num_repeats)) expected_new2old = torch.tensor([6, 1, 0, 3]).to(device) assert torch.all(torch.eq(new2old, expected_new2old)) for device in self.devices: ragged = k2.RaggedInt('[ [1 3] [1 2] [1] [1 4]]').to(device) unique, num_repeats, new2old = k2.ragged.unique_sequences( ragged, need_num_repeats=True, need_new2old_indexes=True) expected = k2.RaggedInt('[[1] [1 2] [1 3] [1 4]]') self.assertEqual(str(unique), str(expected)) expected_num_repeats = k2.RaggedInt('[[1 1 1 1]]') self.assertEqual(str(num_repeats), str(expected_num_repeats)) # CAUTION: The output sublists are ordered by their hash value! expected_new2old = torch.tensor([2, 1, 0, 3]).to(device) assert torch.all(torch.eq(new2old, expected_new2old))
def test_unique_sequences_three_axes(self): for device in self.devices: ragged = k2.RaggedInt( '[ [[1] [1 2] [1 3] [1] [1 3]] [[1 4] [1 2] [1 3] [1 3] [1 2] [1]] ]' # noqa ).to(device) unique, num_repeats, new2old = k2.ragged.unique_sequences( ragged, need_num_repeats=True, need_new2old_indexes=True) expected = k2.RaggedInt( '[ [[1] [1 2] [1 3]] [[1] [1 2] [1 3] [1 4]] ]') self.assertEqual(str(unique), str(expected)) expected_num_repeats = k2.RaggedInt('[ [2 1 2] [1 2 2 1] ]') self.assertEqual(str(num_repeats), str(expected_num_repeats)) expected_new2old = torch.tensor([0, 1, 2, 10, 6, 7, 5]).to(device) assert torch.all(torch.eq(new2old, expected_new2old)) for device in self.devices: ragged = k2.RaggedInt( '[ [[1 3] [1] [1 2]] [[1 2] [1 3] [1 4 5] [1]] ]').to(device) unique, num_repeats, new2old = k2.ragged.unique_sequences( ragged, need_num_repeats=True, need_new2old_indexes=True) expected = k2.RaggedInt( '[ [[1] [1 2] [1 3]] [[1] [1 2] [1 3] [1 4 5]] ]') self.assertEqual(str(unique), str(expected)) expected_num_repeats = k2.RaggedInt('[[1 1 1 ] [1 1 1 1]]') self.assertEqual(str(num_repeats), str(expected_num_repeats)) # CAUTION: The output sublists are ordered by their hash value! expected_new2old = torch.tensor([1, 2, 0, 6, 3, 4, 5]).to(device) assert torch.all(torch.eq(new2old, expected_new2old))
def test_convert_attr_to_ragged(self): for device in self.devices: s = ''' 0 1 1 0.1 1 2 2 0.2 2 3 -1 0.3 3 ''' fsa = k2.Fsa.from_str(s).to(device) fsa.tensor_attr1 = torch.tensor([1, 2, 3, 4, 0, 6], dtype=torch.int32, device=device)[::2] fsa.convert_attr_to_ragged_(name='tensor_attr1', remove_eps=False) expected = k2.RaggedInt('[ [1] [3] [0] ]') assert str(fsa.tensor_attr1) == str(expected) fsa.tensor_attr2 = torch.tensor([1, 0, -1], dtype=torch.int32, device=device) fsa.convert_attr_to_ragged_(name='tensor_attr2', remove_eps=True) expected = k2.RaggedInt('[ [1] [] [-1] ]') assert str(fsa.tensor_attr2) == str(expected)
def test_create_fsa_vec(self): s1 = ''' 0 1 1 0.1 1 2 -1 0.2 2 ''' s2 = ''' 0 1 -1 10 1 ''' fsa1 = k2.Fsa.from_str(s1) fsa1.aux_labels = k2.RaggedInt('[ [1 0 2] [3 5] ]') fsa2 = k2.Fsa.from_str(s2) fsa2.aux_labels = k2.RaggedInt('[ [5 8 9] ]') fsa = k2.create_fsa_vec([fsa1, fsa2]) self.assertEqual(str(fsa.aux_labels), '[ [ 1 0 2 ] [ 3 5 ] [ 5 8 9 ] ]') fsa = k2.Fsa.from_fsas([fsa1, fsa2]) self.assertEqual(str(fsa.aux_labels), '[ [ 1 0 2 ] [ 3 5 ] [ 5 8 9 ] ]')
def test_ragged_int_from_str(self): s = ''' [ [1 2] [3] ] ''' ragged_int = k2.RaggedInt(s) print(ragged_int) assert torch.all(torch.eq(ragged_int.values(), torch.tensor([1, 2, 3]))) assert ragged_int.dim0() == 2 assert torch.all( torch.eq(ragged_int.row_splits(1), torch.tensor([0, 2, 3]))) self.assertEqual([2, 3], ragged_int.tot_sizes())
def test_from_ragged_int_single_fsa(self): for device in self.devices: ragged_int = k2.RaggedInt('[ [10 20] ]').to(device) fsa = k2.linear_fsa(ragged_int) assert fsa.shape == (1, None, None) assert fsa.device == device expected_arcs = torch.tensor([[0, 1, 10], [1, 2, 20], [2, 3, -1]], dtype=torch.int32, device=device) assert torch.all( torch.eq( fsa.arcs.values()[:, :-1], # skip the last field `scores` expected_arcs)) assert torch.all(torch.eq(fsa.scores, torch.zeros_like(fsa.scores)))
def test_remove_values_eq(self): s = ''' [ [1 2 0] [3 0 2] [0 8 0 6 0] [0] ] ''' src = k2.RaggedInt(s) ans = k2.ragged.remove_values_eq(src, 0) self.assertEqual(str(ans), '[ [ 1 2 ] [ 3 2 ] [ 8 6 ] [ ] ]') ans = k2.ragged.remove_values_eq(src, 1) self.assertEqual(str(ans), '[ [ 2 0 ] [ 3 0 2 ] [ 0 8 0 6 0 ] [ 0 ] ]') ans = k2.ragged.remove_values_eq(src, 6) self.assertEqual(str(ans), '[ [ 1 2 0 ] [ 3 0 2 ] [ 0 8 0 0 ] [ 0 ] ]') ans = k2.ragged.remove_values_eq(src, 8) self.assertEqual(str(ans), '[ [ 1 2 0 ] [ 3 0 2 ] [ 0 0 6 0 ] [ 0 ] ]')
def test_remove_values_leq(self): s = ''' [ [1 2 0] [3 0 2] [0 8 0 6 0] [0] ] ''' for device in self.devices: src = k2.RaggedInt(s).to(device) ans = k2.ragged.remove_values_leq(src, 0) self.assertEqual(str(ans), '[ [ 1 2 ] [ 3 2 ] [ 8 6 ] [ ] ]') ans = k2.ragged.remove_values_leq(src, 1) self.assertEqual(str(ans), '[ [ 2 ] [ 3 2 ] [ 8 6 ] [ ] ]') ans = k2.ragged.remove_values_leq(src, 6) self.assertEqual(str(ans), '[ [ ] [ ] [ 8 ] [ ] ]') ans = k2.ragged.remove_values_leq(src, 8) self.assertEqual(str(ans), '[ [ ] [ ] [ ] [ ] ]')
def top_k(self, k: int) -> 'Nbest': '''Get a subset of paths in the Nbest. The resulting Nbest is regular in that each sequence (i.e., utterance) has the same number of paths (k). We select the top-k paths according to the total_scores of each path. If a utterance has less than k paths, then its last path, after sorting by tot_scores in descending order, is repeated so that each utterance has exactly k paths. Args: k: Number of paths in each utterance. Returns: Return a new Nbest with a regular shape. ''' ragged_scores = self.total_scores() # indexes contains idx01's for self.shape # ragged_scores.values()[indexes] is sorted indexes = k2.ragged.sort_sublist(ragged_scores, descending=True, need_new2old_indexes=True) ragged_indexes = k2.RaggedInt(self.shape, indexes) padded_indexes = k2.ragged.pad(ragged_indexes, mode='replicate', value=-1) assert torch.ge(padded_indexes, 0).all(), \ 'Some utterances contain empty ' \ f'n-best: {self.shape.row_splits(1)}' # Select the idx01's of top-k paths of each utterance top_k_indexes = padded_indexes[:, :k].flatten().contiguous() top_k_fsas = k2.index_fsa(self.fsa, top_k_indexes) top_k_shape = k2.ragged.regular_ragged_shape(dim0=self.shape.dim0(), dim1=k) return Nbest(top_k_fsas, top_k_shape)
def test_1d(self): devices = [torch.device('cpu')] if torch.cuda.is_available(): devices.append(torch.device('cuda', 0)) for device in devices: row_splits1 = torch.tensor([0, 3, 5, 6, 6, 9], dtype=torch.int32, device=device) # we don't need to call shape2.to(device) here as shape2 # will be on the same device as row_splits shape2 = k2.create_ragged_shape2(row_splits1, None, 9) values = torch.tensor([1, 0, 4, 2, 3, 0, 4, 5, 2], dtype=torch.int32, device=device) ragged2 = k2.RaggedInt(shape2, values) # contiguous src = torch.tensor([0, 2, 0, 10, 0, -1], dtype=torch.int32, device=device) ans = k2.simple_ragged_index_select(src, ragged2) self.assertEqual(ans.dtype, src.dtype) self.assertEqual(ans.numel(), shape2.dim0()) expected = torch.tensor([2, 10, 0, 0, -1], dtype=torch.int32, device=device) self.assertTrue(torch.allclose(ans, expected)) # non-contiguous src = src.expand(3, -1).t().flatten()[::3] self.assertFalse(src.is_contiguous()) self.assertEqual(src.stride(0), 3) ans = k2.simple_ragged_index_select(src, ragged2) self.assertEqual(ans.dtype, src.dtype) self.assertEqual(ans.numel(), shape2.dim0()) self.assertTrue(ans.is_contiguous()) self.assertEqual(ans.stride(0), 1) expected = torch.tensor([2, 10, 0, 0, -1], dtype=torch.int32, device=device) self.assertTrue(torch.allclose(ans, expected))
def test_from_ragged_int_two_fsas(self): devices = [torch.device('cpu')] if torch.cuda.is_available(): devices.append(torch.device('cuda', 0)) for device in devices: ragged_int = k2.RaggedInt('[ [10 20] [100 200 300] ]').to(device) fsa = k2.linear_fsa(ragged_int) assert fsa.shape == (2, None, None) assert fsa.device == device expected_arcs = torch.tensor( [[0, 1, 10], [1, 2, 20], [2, 3, -1], [0, 1, 100], [1, 2, 200], [2, 3, 300], [3, 4, -1]], dtype=torch.int32, device=device) assert torch.all( torch.eq( fsa.arcs.values()[:, :-1], # skip the last field `scores` expected_arcs)) assert torch.all(torch.eq(fsa.scores, torch.zeros_like(fsa.scores)))
def test_pad(self): s = ''' [ [ 1 2 ] [ 3 ] [ ] [ 4 5 6 ] [ 7 8 9 10 ] ] ''' for device in self.devices: src = k2.RaggedInt(s).to(device) value = random.randint(0, 1000) ans = k2.ragged.pad(src, 'constant', value) expected = torch.ones( (5, 4), dtype=torch.int32, device=device) * value expected[0, 0] = 1 expected[0, 1] = 2 expected[1, 0] = 3 expected[3, 0] = 4 expected[3, 1] = 5 expected[3, 2] = 6 expected[4, 0] = 7 expected[4, 1] = 8 expected[4, 2] = 9 expected[4, 3] = 10 assert torch.all(torch.eq(ans, expected))
def test(self): for device in self.devices: src = torch.tensor([1, 2, 3, 4, 5, 6, 7], dtype=torch.int32, device=device) index_row_splits = torch.tensor([0, 2, 2, 3, 7], dtype=torch.int32, device=device) index_shape = k2.ragged.create_ragged_shape2( index_row_splits, None, 7) index_values = torch.tensor([0, 3, 2, 3, 5, 1, 3], dtype=torch.int32, device=device) ragged_index = k2.RaggedInt(index_shape, index_values) ans = k2.index(src, ragged_index) self.assertTrue(torch.allclose(ans.row_splits(1), index_row_splits)) expected_values = torch.tensor([1, 4, 3, 4, 6, 2, 4], dtype=torch.int32, device=device) self.assertTrue(torch.allclose(ans.values(), expected_values))
def test_with_negative_1(self): for device in self.devices: src = torch.tensor([0, 1, 2, 3], dtype=torch.float32, requires_grad=True, device=device) indexes = k2.RaggedInt( '[ [1 2 -1] [0 3] [-1] [0 2 3 1 3] [] ]').to(device) ans = k2.index_and_sum(src, indexes) expected = torch.tensor([1 + 2, 0 + 3, 0, 0 + 2 + 3 + 1 + 3, 0]).to(src) assert torch.allclose(ans, expected) # now for autograd scale = torch.tensor([10, 20, 30, 40, 50]).to(device) (ans * scale).sum().backward() expected_grad = torch.empty_like(src.grad) expected_grad[0] = scale[1] + scale[3] expected_grad[1] = scale[0] + scale[3] expected_grad[2] = scale[0] + scale[3] expected_grad[3] = scale[1] + scale[3] * 2 assert torch.allclose(src.grad, expected_grad)
def test_aux_as_ragged(self): s = ''' 0 1 1 0 0 1 0 0 0 3 2 0 1 2 3 0 1 3 4 0 2 1 5 0 2 5 -1 0 3 1 6 0 4 5 -1 0 5 ''' fsa = k2.Fsa.from_str(s) assert fsa.device.type == 'cpu' aux_row_splits = torch.tensor([0, 2, 3, 3, 6, 6, 7, 8, 10, 11], dtype=torch.int32) aux_shape = k2.ragged.create_ragged_shape2(aux_row_splits, None, 11) aux_values = torch.tensor([1, 2, 3, 5, 6, 7, 8, -1, 9, 10, -1], dtype=torch.int32) fsa.aux_labels = k2.RaggedInt(aux_shape, aux_values) dest = k2.invert(fsa) print(dest) # will print aux_labels as well
def test(self): devices = [torch.device('cpu')] if torch.cuda.is_available(): devices.append(torch.device('cuda', 0)) for device in devices: src = torch.tensor([1, 2, 3, 4, 5, 6, 7], dtype=torch.int32, device=device) index_row_splits = torch.tensor([0, 2, 2, 3, 7], dtype=torch.int32, device=device) index_shape = k2.create_ragged_shape2(index_row_splits, None, 7) index_values = torch.tensor([0, 3, 2, 3, 5, 1, 3], dtype=torch.int32, device=device) ragged_index = k2.RaggedInt(index_shape, index_values) ans = k2.index_tensor_with_ragged_int(src, ragged_index) self.assertTrue(torch.allclose(ans.row_splits(1), index_row_splits)) expected_values = torch.tensor([1, 4, 3, 4, 6, 2, 4], dtype=torch.int32, device=device) self.assertTrue(torch.allclose(ans.values(), expected_values))
def test_without_empty_list(self): devices = [torch.device('cpu')] if torch.cuda.is_available(): devices.append(torch.device('cuda', 0)) for device in devices: s = ''' 0 1 0 0 0 1 1 0 1 2 -1 0 2 ''' scores = torch.tensor([1, 2, 3], dtype=torch.float32, device=device, requires_grad=True) scores_copy = scores.detach().clone().requires_grad_(True) src = k2.Fsa.from_str(s).to(device) src.scores = scores src.attr1 = "hello" src.attr2 = "k2" float_attr = torch.tensor([0.1, 0.2, 0.3], dtype=torch.float32, requires_grad=True, device=device) src.float_attr = float_attr.detach().clone().requires_grad_(True) src.int_attr = torch.tensor([1, 2, 3], dtype=torch.int32, device=device) src.ragged_attr = k2.RaggedInt( '[ [10 20] [30 40 50] [60 70] ]').to(device) ragged_arc, arc_map = _k2.remove_epsilon(src.arcs, src.properties) dest = k2.utils.fsa_from_unary_function_ragged( src, ragged_arc, arc_map) assert dest.attr1 == src.attr1 assert dest.attr2 == src.attr2 if device.type == 'cpu': expected_arc_map = k2.RaggedInt('[ [1] [0 2] [2] ]') self.assertEqual(str(arc_map), str(expected_arc_map)) expected_int_attr = k2.RaggedInt('[ [2] [1 3] [3] ]') self.assertEqual(str(dest.int_attr), str(expected_int_attr)) expected_ragged_attr = k2.RaggedInt( '[ [30 40 50] [10 20 60 70] [60 70] ]') self.assertEqual(str(dest.ragged_attr), str(expected_ragged_attr)) expected_float_attr = torch.empty_like(dest.float_attr) expected_float_attr[0] = float_attr[1] expected_float_attr[1] = float_attr[0] + float_attr[2] expected_float_attr[2] = float_attr[2] assert torch.all(torch.eq(dest.float_attr, expected_float_attr)) expected_scores = torch.empty_like(dest.scores) expected_scores[0] = scores_copy[1] expected_scores[1] = scores_copy[0] + scores_copy[2] expected_scores[2] = scores_copy[2] assert torch.all(torch.eq(dest.scores, expected_scores)) else: # For CUDA, the arc_map differs, but the resulting FSA # is equivalent. expected_arc_map = k2.RaggedInt('[ [0 2] [1] [2] ]') self.assertEqual(str(arc_map), str(expected_arc_map)) expected_int_attr = k2.RaggedInt('[ [1 3] [2] [3] ]') self.assertEqual(str(dest.int_attr), str(expected_int_attr)) expected_ragged_attr = k2.RaggedInt( '[ [10 20 60 70] [30 40 50] [60 70] ]') self.assertEqual(str(dest.ragged_attr), str(expected_ragged_attr)) expected_float_attr = torch.empty_like(dest.float_attr) expected_float_attr[0] = float_attr[0] + float_attr[2] expected_float_attr[1] = float_attr[1] expected_float_attr[2] = float_attr[2] assert torch.all(torch.eq(dest.float_attr, expected_float_attr)) expected_scores = torch.empty_like(dest.scores) expected_scores[0] = scores_copy[0] + scores_copy[2] expected_scores[1] = scores_copy[1] expected_scores[2] = scores_copy[2] assert torch.all(torch.eq(dest.scores, expected_scores)) scale = torch.tensor([10, 20, 30]).to(float_attr) (dest.float_attr * scale).sum().backward() (expected_float_attr * scale).sum().backward() assert torch.all(torch.eq(src.float_attr.grad, float_attr.grad)) (dest.scores * scale).sum().backward() (expected_scores * scale).sum().backward() assert torch.all(torch.eq(scores.grad, scores_copy.grad))
def test_append_axis1(self): ragged1 = k2.RaggedInt('[ [1 2 3] [] [4 5] ]') ragged2 = k2.RaggedInt('[ [10 20] [8] [9 10] ]') ragged = k2.ragged.append([ragged1, ragged2], axis=1) self.assertEqual(str(ragged), '[ [ 1 2 3 10 20 ] [ 8 ] [ 4 5 9 10 ] ]')
def test(self): for device in self.devices: s = ''' 0 1 2 10 0 1 1 20 1 2 -1 30 2 ''' src = k2.Fsa.from_str(s).to(device).requires_grad_(True) src.float_attr = torch.tensor([0.1, 0.2, 0.3], dtype=torch.float32, requires_grad=True, device=device) src.int_attr = torch.tensor([1, 2, 3], dtype=torch.int32, device=device) src.ragged_attr = k2.RaggedInt('[[1 2 3] [5 6] []]').to(device) src.attr1 = 'src' src.attr2 = 'fsa' ragged_arc, arc_map = _k2.arc_sort(src.arcs, need_arc_map=True) dest = k2.utils.fsa_from_unary_function_tensor( src, ragged_arc, arc_map) assert torch.allclose( dest.float_attr, torch.tensor([0.2, 0.1, 0.3], dtype=torch.float32, device=device)) assert torch.all( torch.eq( dest.scores, torch.tensor([20, 10, 30], dtype=torch.float32, device=device))) assert torch.all( torch.eq( dest.int_attr, torch.tensor([2, 1, 3], dtype=torch.int32, device=device))) expected_ragged_attr = k2.RaggedInt('[ [5 6] [1 2 3] []]') self.assertEqual(str(dest.ragged_attr), str(expected_ragged_attr)) assert dest.attr1 == src.attr1 assert dest.attr2 == src.attr2 # now for autograd scale = torch.tensor([10, 20, 30], device=device) (dest.float_attr * scale).sum().backward() (dest.scores * scale).sum().backward() expected_grad = torch.tensor([20, 10, 30], dtype=torch.float32, device=device) assert torch.all(torch.eq(src.float_attr.grad, expected_grad)) assert torch.all(torch.eq(src.scores.grad, expected_grad))
def test_with_negative_1(self): devices = [torch.device('cpu')] if torch.cuda.is_available(): devices.append(torch.device('cuda', 0)) for device in devices: s = ''' 0 1 2 10 0 1 1 20 1 2 -1 30 2 ''' src = k2.Fsa.from_str(s).to(device).requires_grad_(True) src.float_attr = torch.tensor([0.1, 0.2, 0.3], dtype=torch.float32, requires_grad=True, device=device) src.int_attr = torch.tensor([1, 2, 3], dtype=torch.int32, device=device) src.ragged_attr = k2.RaggedInt('[[1 2 3] [5 6] []]').to(device) src.attr1 = 'src' src.attr2 = 'fsa' ragged_arc, arc_map = _k2.add_epsilon_self_loops(src.arcs, need_arc_map=True) dest = k2.utils.fsa_from_unary_function_tensor( src, ragged_arc, arc_map) assert torch.allclose( dest.float_attr, torch.tensor([0.0, 0.1, 0.2, 0.0, 0.3], dtype=torch.float32, device=device)) assert torch.all( torch.eq( dest.scores, torch.tensor([0, 10, 20, 0, 30], dtype=torch.float32, device=device))) assert torch.all( torch.eq( dest.int_attr, torch.tensor([0, 1, 2, 0, 3], dtype=torch.int32, device=device))) expected_ragged_attr = k2.RaggedInt('[ [] [1 2 3] [5 6] [] []]') self.assertEqual(str(dest.ragged_attr), str(expected_ragged_attr)) assert dest.attr1 == src.attr1 assert dest.attr2 == src.attr2 # now for autograd scale = torch.tensor([10, 20, 30, 40, 50], device=device) (dest.float_attr * scale).sum().backward() (dest.scores * scale).sum().backward() expected_grad = torch.tensor([20, 30, 50], dtype=torch.float32, device=device) assert torch.all(torch.eq(src.float_attr.grad, expected_grad)) assert torch.all(torch.eq(src.scores.grad, expected_grad))
def test_with_empty_list(self): for device in self.devices: s = ''' 0 1 0 0 0 1 1 0 1 2 -1 0 2 ''' scores = torch.tensor([1, 2, 3], dtype=torch.float32, device=device, requires_grad=True) scores_copy = scores.detach().clone().requires_grad_(True) src = k2.Fsa.from_str(s).to(device) src.scores = scores src.attr1 = "hello" src.attr2 = "k2" float_attr = torch.tensor([0.1, 0.2, 0.3], dtype=torch.float32, requires_grad=True, device=device) src.float_attr = float_attr.detach().clone().requires_grad_(True) src.int_attr = torch.tensor([1, 2, 3], dtype=torch.int32, device=device) src.ragged_attr = k2.RaggedInt( '[ [10 20] [30 40 50] [60 70] ]').to(device) ragged_arc, arc_map = _k2.remove_epsilon_and_add_self_loops( src.arcs, src.properties) dest = k2.utils.fsa_from_unary_function_ragged( src, ragged_arc, arc_map) assert dest.attr1 == src.attr1 assert dest.attr2 == src.attr2 expected_arc_map = k2.RaggedInt('[ [] [1] [0 2] [] [2] ]') self.assertEqual(str(arc_map), str(expected_arc_map)) expected_int_attr = k2.RaggedInt('[ [] [2] [1 3] [] [3] ]') self.assertEqual(str(dest.int_attr), str(expected_int_attr)) expected_ragged_attr = k2.RaggedInt( '[ [] [30 40 50] [10 20 60 70] [] [60 70] ]') self.assertEqual(str(dest.ragged_attr), str(expected_ragged_attr)) expected_float_attr = torch.empty_like(dest.float_attr) expected_float_attr[0] = 0 expected_float_attr[1] = float_attr[1] expected_float_attr[2] = float_attr[0] + float_attr[2] expected_float_attr[3] = 0 expected_float_attr[4] = float_attr[2] assert torch.all(torch.eq(dest.float_attr, expected_float_attr)) expected_scores = torch.empty_like(dest.scores) expected_scores[0] = 0 expected_scores[1] = scores_copy[1] expected_scores[2] = scores_copy[0] + scores_copy[2] expected_scores[3] = 0 expected_scores[4] = scores_copy[2] assert torch.all(torch.eq(dest.scores, expected_scores)) scale = torch.tensor([10, 20, 30, 40, 50]).to(float_attr) (dest.float_attr * scale).sum().backward() (expected_float_attr * scale).sum().backward() assert torch.all(torch.eq(src.float_attr.grad, float_attr.grad)) (dest.scores * scale).sum().backward() (expected_scores * scale).sum().backward() assert torch.all(torch.eq(scores.grad, scores_copy.grad))
def test(self): devices = [torch.device('cpu')] if torch.cuda.is_available() and k2.with_cuda: devices.append(torch.device('cuda', 0)) for device in devices: for need_map in [True, False]: s = ''' 0 1 2 10 0 1 1 20 1 2 -1 30 2 ''' src = k2.Fsa.from_str(s).to(device).requires_grad_(True) src.float_attr = torch.tensor([0.1, 0.2, 0.3], dtype=torch.float32, requires_grad=True, device=device) src.int_attr = torch.tensor([1, 2, 3], dtype=torch.int32, device=device) src.ragged_attr = k2.RaggedInt('[[1 2 3] [5 6] []]').to(device) src.attr1 = 'src' src.attr2 = 'fsa' if need_map: dest, arc_map = k2.expand_ragged_attributes( src, ret_arc_map=True) else: dest = k2.expand_ragged_attributes(src) assert torch.allclose( dest.float_attr, torch.tensor([0.1, 0.2, 0.0, 0.0, 0.0, 0.3], dtype=torch.float32, device=device)) assert torch.all( torch.eq( dest.scores, torch.tensor([10, 20, 0, 0, 0, 30], dtype=torch.float32, device=device))) assert torch.all( torch.eq( dest.int_attr, torch.tensor([1, 2, 0, 0, 0, 3], dtype=torch.int32, device=device))) assert torch.all( torch.eq( dest.ragged_attr, torch.tensor([1, 5, 2, 3, 6, -1], dtype=torch.float32, device=device))) # non-tensor attributes... assert dest.attr1 == src.attr1 assert dest.attr2 == src.attr2 # now for autograd scale = torch.tensor([10, 20, 10, 10, 10, 30], device=device) (dest.float_attr * scale).sum().backward() (dest.scores * scale).sum().backward() expected_grad = torch.tensor([10, 20, 30], dtype=torch.float32, device=device) assert torch.all(torch.eq(src.float_attr.grad, expected_grad)) assert torch.all(torch.eq(src.scores.grad, expected_grad))
import k2 s = ''' 0 1 2 0.1 1 2 -1 0.2 2 ''' fsa = k2.Fsa.from_str(s) fsa.aux_labels = k2.RaggedInt('[ [10 20] [-1] ]') inverted_fsa = k2.invert(fsa) fsa.draw('before_invert_aux.svg', title='before invert with ragged tensors as aux_labels') inverted_fsa.draw('after_invert_aux.svg', title='after invert')