[WIP]: Move k2.Fsa to C++

[csukuangfj]

This is an initial attempt to move most of the code, if not all, about k2.Fsa to C++.

@danpovey I am not sure if this is the right direction to go.

My intention is to make k2 a pure C++ library and what we provide are just some *.so libraries.
(Of course, it does not change the way to install k2)

Everything currently written in Python can be moved to C++, I believe (though it needs some time to verify whether it is feasible).

The advantage is that it makes k2 more production-friendly.

---

The following is some example code for demonstration of the current commit in this pull-request.
(Note: k2.Fsa.arc_sort() supports autograd)

#!/usr/bin/env python3

import torch
from _k2.ragged import Fsa

s = """
0 1 2 10 0.2
0 1 1 20 0.3
1 2 -1 30 4.
2
"""
fsa = Fsa(s, ["aux_labels"])
fsa.requires_grad_(True)
print("fsa.scores:", fsa.scores)

sorted_fsa = fsa.arc_sort()
print("sorted_fsa.scores:", sorted_fsa.scores)
(sorted_fsa.scores * torch.tensor([10, 20, 30])).sum().backward()
print(fsa.grad)

print(fsa)
print("-" * 10, "after sort", "-" * 10)
print(sorted_fsa)

The output of the above code is:

fsa.scores: tensor([0.2000, 0.3000, 4.0000], requires_grad=True)
sorted_fsa.scores: tensor([0.3000, 0.2000, 4.0000], grad_fn=<ArcSortFunction>>)
tensor([20., 10., 30.])
k2.Fsa: 0 1 2 10 0.2
0 1 1 20 0.3
1 2 -1 30 4
2

---------- after sort ----------
k2.Fsa: 0 1 1 0.3
0 1 2 0.2
1 2 -1 4
2

[csukuangfj]

This is how we do backprop for ArcSort.

[danpovey]

Ah OK, cool.

[danpovey]

I'm not sure that it is is a good idea to call this simply Tensor, IMO there is too much pontential for confusion here(?) RaggedTensor might be a possibility? Or just Ragged, maybe?

[csukuangfj]

The Tensor is put inside the module k2.ragged; its fully qualified name is k2.ragged.Tensor.

If it is called RaggedTensor, shall we put it in k2, or k2.ragged? That is, shall we call it
k2.RaggedTensor or k2.ragged.RaggedTensor?

By the way, this file is going to be deleted as the documentation is now contained in the C++ header files.
This pull request is based on #812 .

If you like the idea (i.e., moving all possible parts from Python to C++), I will close this pullrequest and
add commits from the pull-request to #812

[danpovey]

I think we could maybe keep it as k2.ragged.RaggedTensor but also import it into k2.
I still have to think about this. I'm impressed that you figured out how to do it, and we probably need to do this at
some point, but for me the main question in my mind is what is going to happen to class k2.Fsa. That has a bunch of
logic, e.g. about attributes, that I wonder how it will be transferred to C++; and once it is C++, who will be able to
maintain it!

[danpovey]

I'm curious how we access arrays of type Arc from Python.

[csukuangfj]

That is doable, like what we are currently doing. That is, convert Array1<Arc> to a 2-d int32 torch.Tensor, though the last column of the tensor actually contains floats, but is reinterpreted as ints.

[csukuangfj]

#814 (comment)

but for me the main question in my mind is what is going to happen to class k2.Fsa. That has a bunch of
logic, e.g. about attributes, that I wonder how it will be transferred to C++; and once it is C++, who will be able to
maintain it!

I would propose to let these three maps manage all attributes of k2.Fsa.

We can add

• __getattr__
• __setattr__
  to RaggedArc when wrapping it to Python. In this way, we can support assigning
  arbitrary attributes to an Fsa in Python.
• If the attribute is a tensor, we save it to tensor_attrs.
• If the attribute is a ragged tensor, we save it to ragged_tensor_attrs.
• Otherwise, we save it to other_attrs.

all_attr_names is to ensure that an attribute has exactly only one type.

[danpovey]

Ah, I see, I missed this.
I feel like the name RaggedArc is misleading, since this implements the full Fsa logic.
Maybe something like FsaWrapper or FsaClass would be clearer.
My main concern with this, I suppose, is that the code becomes less easy for others to understand,
view or maintain, because C++ is less accessible. But it does solve a potential problem with deployment
in a C++-only environment.
Also, if we go down this road, it's going to be a lot of work, and I wonder whether it can be broken up
into stages somehow? E.g. first adding the Any stuff, and later move class Fsa to C++?

[csukuangfj]

Also, if we go down this road, it's going to be a lot of work, and I wonder whether it can be broken up
into stages somehow? E.g. first adding the Any stuff, and later move class Fsa to C++?

Yes, I agree. I will try to get RaggedAny done first.

[csukuangfj]

My main concern with this, I suppose, is that the code becomes less easy for others to understand,

I think this can be mitigated by detailed documentation.

[pkufool]

Sounds exciting!
@csukuangfj You could submit RaggedAny first, I'll take part in some stuff of moving class Fsa to C++, if possible.

[csukuangfj]

@pkufool
Thanks!

[csukuangfj]

Shall we move this to a separate branch, say v2.0-pre?

[danpovey]

Yes, that's a good idea I think.

[pkufool]

Since this is not urgent, I want to have a try to take over this task, doing it gradually (about 30% of my time), finally finish this in about one or two month. @csukuangfj What do you think.

[csukuangfj]

Since this is not urgent, I want to have a try to take over this task, doing it gradually (about 30% of my time), finally finish this in about one or two months. @csukuangfj What do you think.

👍
You can go ahead and I will help you to do this.

[pkufool]

Since this is not urgent, I want to have a try to take over this task, doing it gradually (about 30% of my time), finally finish this in about one or two months. @csukuangfj What do you think.

👍
You can go ahead and I will help you to do this.

Thanks! of course, I need your help!

[csukuangfj]

@danpovey

This shows how we access attributes of FSAs from Python, whose
underlying implementation is in C++.

Will show how to do autograd on these attributes for FSA operations in later commits.

[csukuangfj]

There are several TODOs:

• fsa_from_unary_function_tensor
  

  k2/k2/python/k2/utils.py

  Line 423 in 8030001

  def fsa_from_unary_function_tensor(src: Fsa, dest_arcs: _k2.RaggedArc,

• fsa_from_unary_function_ragged
  

  k2/k2/python/k2/utils.py

  Line 467 in 8030001

  def fsa_from_unary_function_ragged(src: Fsa,

• fsa_from_binary_function_tensor
  

  k2/k2/python/k2/utils.py

  Line 552 in 8030001

  def fsa_from_binary_function_tensor(a_fsa: Fsa, b_fsa: Fsa,

• _IndexSelectFunction
  

  k2/k2/python/k2/ops.py

  Line 27 in 8030001

  class _IndexSelectFunction(torch.autograd.Function):

After the above items are done, we can start to move fsa_algo.py to C++.

@pkufool You can select which item(s) you want to implement and I will take the remaining.

[pkufool]

@csukuangfj I'll try the first two first.

[csukuangfj]

This shows how to implement autograd for fsa.get_forward_scores().

[csukuangfj]

Merging it to the branch v2.0-pre and we (@pkufool ) will add more commits incrementally.