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')
