Wrap Fsa to Python. (k2-fsa#193)

* wrap Fsa to Python.

* add more documentation.

* resolve some comments.

* fix style issues.

* resolve some comments.

* Make Fsa a Python class.

Now the API is more flexible and is much simpler.

* fix style issues.

* fix a typo.

* add more documentation.

* fix style issues.

* wrap Array2<T> to Python.

* fix style issues.

* Wrap DenseFsaVec to Python.

* add visualization support for FSA.

Author: csukuangfj

k2/csrc/CMakeLists.txt

@@ -76,6 +76,6 @@ function(k2_add_cuda_test name)
   )
 endfunction()
 
-foreach (name IN LISTS cuda_tests)
+foreach(name IN LISTS cuda_tests)
   k2_add_cuda_test(${name})
-endforeach ()
+endforeach()

k2/csrc/array_ops_inl.h

@@ -241,7 +241,7 @@ Array1<T> Append(int32_t num_arrays, const Array1<T> **src) {
   std::vector<int32_t> row_splits_vec(num_arrays + 1);
   int32_t sum = 0, max_dim = 0;
   row_splits_vec[0] = sum;
-  for (int32_t i = 0; i < num_arrays; i++) {
+  for (int32_t i = 0; i < num_arrays; ++i) {
     int32_t dim = src[i]->Dim();
     if (dim > max_dim) max_dim = dim;
     sum += dim;
@@ -256,7 +256,7 @@ Array1<T> Append(int32_t num_arrays, const Array1<T> **src) {
     // a simple loop is faster, although the other branches should still work on
     // CPU.
     int32_t elem_size = src[0]->ElementSize();
-    for (int32_t i = 0; i < num_arrays; i++) {
+    for (int32_t i = 0; i < num_arrays; ++i) {
       int32_t this_dim = src[i]->Dim();
       const T *this_src_data = src[i]->Data();
       memcpy(static_cast<void *>(ans_data),
@@ -268,7 +268,7 @@ Array1<T> Append(int32_t num_arrays, const Array1<T> **src) {
     Array1<int32_t> row_splits(c, row_splits_vec);
     const int32_t *row_splits_data = row_splits.Data();
     std::vector<const T *> src_ptrs_vec(num_arrays);
-    for (int32_t i = 0; i < num_arrays; i++) src_ptrs_vec[i] = src[i]->Data();
+    for (int32_t i = 0; i < num_arrays; ++i) src_ptrs_vec[i] = src[i]->Data();
     Array1<const T *> src_ptrs(c, src_ptrs_vec);
     const T **src_ptrs_data = src_ptrs.Data();
     int32_t avg_input_size = ans_size / num_arrays;
@@ -305,10 +305,10 @@ Array1<T> Append(int32_t num_arrays, const Array1<T> **src) {
       // them on CPU.
       std::vector<uint64_t> index_map;
       index_map.reserve((2 * ans_size) / block_dim);
-      for (int32_t i = 0; i < num_arrays; i++) {
+      for (int32_t i = 0; i < num_arrays; ++i) {
         int32_t this_array_size = src[i]->Dim();
         int32_t this_num_blocks = NumBlocks(this_array_size, block_dim);
-        for (int32_t j = 0; j < this_num_blocks; j++) {
+        for (int32_t j = 0; j < this_num_blocks; ++j) {
           index_map.push_back((static_cast<uint64_t>(j) << 32) +
                               static_cast<uint64_t>(i));
         }
@@ -362,10 +362,10 @@ void ApplyOpPerSublist(Ragged<T> &src, T default_value, Array1<T> *dst) {
 
   if (c->GetDeviceType() == kCpu) {
     int32_t j = row_splits[0];
-    for (int32_t i = 0; i < num_rows; i++) {
+    for (int32_t i = 0; i < num_rows; ++i) {
       T val = default_value;
       int32_t row_end = row_splits[i + 1];
-      for (; j < row_end; j++) {
+      for (; j < row_end; ++j) {
         T elem = values_data[j];
         val = op(elem, val);
       }

k2/csrc/dtype.h

@@ -18,11 +18,11 @@
 
 namespace k2 {
 
-enum BaseType {      // BaseType is the *general type*
-  kUnknownBase = 0,  // e.g. can use this for structs
-  kFloatBase = 1,
-  kIntBase = 2,   // signed int
-  kUintBase = 3,  // unsigned int
+enum BaseType : int8_t {  // BaseType is the *general type*
+  kUnknownBase = 0,       // e.g. can use this for structs
+  kFloatBase = 1,         // real numbers, e.g., float or double
+  kIntBase = 2,           // signed int, e.g., int8_t, int32_t
+  kUintBase = 3,          // unsigned int, e.g, uint32_t, uint64_t
 };
 
 class DtypeTraits {

k2/csrc/fsa_utils.cu

@@ -1,5 +1,7 @@
 /**
- * @brief Utilities for reading, writing and creating FSAs.
+ * @brief Utilities for creating FSAs.
+ *
+ * Note that serializations are done in Python.
  *
  * @copyright
  * Copyright (c)  2020  Mobvoi Inc.        (authors: Fangjun Kuang)

k2/csrc/fsa_utils.h

@@ -1,5 +1,7 @@
 /**
- * @brief Utilities for reading, writing and creating FSAs.
+ * @brief Utilities for creating FSAs.
+ *
+ * Note that serializations are done in Python.
  *
  * @copyright
  * Copyright (c)  2020  Mobvoi Inc.        (authors: Fangjun Kuang)
@@ -45,6 +47,8 @@ namespace k2 {
 
   CAUTION: We assume that `final_state` has the largest state number.
 
+  CAUTION: The first column has to be in non-decreasing order.
+
   @param [in]   s   The input string. See the above description for its format.
   @param [in]   negate_scores
                     If true, the string form has the weights as costs,

k2/csrc/ragged.h

@@ -462,7 +462,7 @@ struct Ragged {
     return Ragged<T>(new_shape, values);
   }
 
-  Ragged<T> To(ContextPtr ctx) {
+  Ragged<T> To(ContextPtr ctx) const {
     RaggedShape new_shape = shape.To(ctx);
     Array1<T> new_values = values.To(ctx);
     return Ragged<T>(new_shape, new_values);

k2/python/csrc/torch.cu

@@ -12,9 +12,17 @@
 
 #if defined(K2_USE_PYTORCH)
 
+#include "k2/python/csrc/torch/arc.h"
 #include "k2/python/csrc/torch/array.h"
+#include "k2/python/csrc/torch/fsa.h"
+#include "k2/python/csrc/torch/ragged.h"
 
-void PybindTorch(py::module &m) { PybindArray(m); }
+void PybindTorch(py::module &m) {
+  PybindArc(m);
+  PybindArray(m);
+  PybindRagged(m);
+  PybindFsa(m);
+}
 
 #else
 

k2/python/csrc/torch/CMakeLists.txt

@@ -1,6 +1,9 @@
 # please keep the list sorted
 set(torch_srcs
+  arc.cu
   array.cu
+  fsa.cu
+  ragged.cu
   torch_util.cu
 )
 

k2/python/csrc/torch/arc.cu

@@ -0,0 +1,62 @@
+/**
+ * @brief python wrappers for Arc.
+ *
+ * @copyright
+ * Copyright (c)  2020  Mobvoi Inc.        (authors: Fangjun Kuang)
+ *
+ * @copyright
+ * See LICENSE for clarification regarding multiple authors
+ */
+
+#include <string>
+
+#include "k2/csrc/fsa.h"
+#include "k2/python/csrc/torch/arc.h"
+#include "k2/python/csrc/torch/torch_util.h"
+#include "torch/extension.h"
+
+namespace k2 {
+
+static void PybindArcImpl(py::module &m) {
+  using PyClass = Arc;
+  py::class_<PyClass> pyclass(m, "Arc");
+  pyclass.def(py::init<>());
+  pyclass.def(py::init<int32_t, int32_t, int32_t, float>(),
+              py::arg("src_state"), py::arg("dest_state"), py::arg("symbol"),
+              py::arg("score"));
+
+  pyclass.def_readwrite("src_state", &PyClass::src_state)
+      .def_readwrite("dest_state", &PyClass::dest_state)
+      .def_readwrite("symbol", &PyClass::symbol)
+      .def_readwrite("score", &PyClass::score);
+
+  pyclass.def("__str__", [](const PyClass &self) -> std::string {
+    std::ostringstream os;
+    os << self;
+    return os.str();
+  });
+
+  m.def("_float_as_int",
+        [](float f) -> int32_t { return *reinterpret_cast<int32_t *>(&f); });
+
+  m.def("_int_as_float",
+        [](int32_t i) -> float { return *reinterpret_cast<float *>(&i); });
+
+  m.def("_as_int", [](torch::Tensor tensor) -> torch::Tensor {
+    auto scalar_type = ToScalarType<int32_t>::value;
+    return torch::from_blob(
+        tensor.data_ptr(), tensor.sizes(), tensor.strides(),
+        [tensor](void *p) {}, tensor.options().dtype(scalar_type));
+  });
+
+  m.def("_as_float", [](torch::Tensor tensor) -> torch::Tensor {
+    auto scalar_type = ToScalarType<float>::value;
+    return torch::from_blob(
+        tensor.data_ptr(), tensor.sizes(), tensor.strides(),
+        [tensor](void *p) {}, tensor.options().dtype(scalar_type));
+  });
+}
+
+}  // namespace k2
+
+void PybindArc(py::module &m) { k2::PybindArcImpl(m); }

k2/python/csrc/torch/arc.h

@@ -0,0 +1,18 @@
+/**
+ * @brief python wrappers for Arc.
+ *
+ * @copyright
+ * Copyright (c)  2020  Mobvoi Inc.        (authors: Fangjun Kuang)
+ *
+ * @copyright
+ * See LICENSE for clarification regarding multiple authors
+ */
+
+#ifndef K2_PYTHON_CSRC_TORCH_ARC_H_
+#define K2_PYTHON_CSRC_TORCH_ARC_H_
+
+#include "k2/python/csrc/k2.h"
+
+void PybindArc(py::module &m);
+
+#endif  // K2_PYTHON_CSRC_TORCH_ARC_H_

k2/python/csrc/torch/array.cu

@@ -9,26 +9,54 @@
  */
 
 #include <type_traits>
+#include <vector>
 
 #include "c10/core/ScalarType.h"
 #include "k2/csrc/array.h"
+#include "k2/csrc/fsa.h"
 #include "k2/csrc/pytorch_context.h"
 #include "k2/python/csrc/torch/array.h"
 #include "k2/python/csrc/torch/torch_util.h"
 #include "torch/extension.h"
 
 namespace k2 {
 
+template <typename T>
+static void PybindArray2Tpl(py::module &m, const char *name) {
+  using PyClass = Array2<T>;
+  py::class_<PyClass> pyclass(m, name);
+  pyclass.def("tensor",
+              [](PyClass &self) -> torch::Tensor { return ToTensor(self); });
+
+  pyclass.def_static(
+      "from_tensor",
+      [](torch::Tensor &tensor) -> PyClass {
+        return FromTensor<T>(tensor, Array2Tag{});
+      },
+      py::arg("tensor"));
+
+  // the following functions are for testing only
+  pyclass.def(
+      "get", [](PyClass &self, int32_t i) -> Array1<T> { return self[i]; },
+      py::arg("i"));
+
+  pyclass.def("__str__", [](const PyClass &self) {
+    std::ostringstream os;
+    os << self;
+    return os.str();
+  });
+}
+
 template <typename T>
 static void PybindArray1Tpl(py::module &m, const char *name) {
   using PyClass = Array1<T>;
   py::class_<PyClass> pyclass(m, name);
-  pyclass.def(py::init<>());
-  pyclass.def("tensor", [](PyClass &self) { return ToTensor(self); });
+  pyclass.def("tensor",
+              [](PyClass &self) -> torch::Tensor { return ToTensor(self); });
 
   pyclass.def_static(
       "from_tensor",
-      [](torch::Tensor &tensor) { return FromTensor<T>(tensor); },
+      [](torch::Tensor &tensor) -> PyClass { return FromTensor<T>(tensor); },
       py::arg("tensor"));
 
   // the following functions are for testing only
@@ -46,30 +74,66 @@ static void PybindArrayImpl(py::module &m) {
   // users should not use classes with prefix `_` in Python.
   PybindArray1Tpl<float>(m, "_FloatArray1");
   PybindArray1Tpl<int>(m, "_Int32Array1");
+  PybindArray1Tpl<Arc>(m, "_ArcArray1");
+
+  PybindArray2Tpl<float>(m, "_FloatArray2");
+  PybindArray2Tpl<int>(m, "_Int32Array2");
 
   // the following functions are for testing purposes
   // and they can be removed later.
-  m.def("get_cpu_float_array1", []() {
+  m.def("get_cpu_float_array1", []() -> Array1<float> {
     return Array1<float>(GetCpuContext(), {1, 2, 3, 4});
   });
 
-  m.def("get_cpu_int_array1", []() {
+  m.def("get_cpu_int_array1", []() -> Array1<int32_t> {
     return Array1<int32_t>(GetCpuContext(), {1, 2, 3, 4});
   });
 
   m.def(
       "get_cuda_float_array1",
-      [](int32_t gpu_id = -1) {
+      [](int32_t gpu_id = -1) -> Array1<float> {
         return Array1<float>(GetCudaContext(gpu_id), {0, 1, 2, 3});
       },
       py::arg("gpu_id") = -1);
 
   m.def(
       "get_cuda_int_array1",
-      [](int32_t gpu_id = -1) {
+      [](int32_t gpu_id = -1) -> Array1<int32_t> {
         return Array1<int32_t>(GetCudaContext(gpu_id), {0, 1, 2, 3});
       },
       py::arg("gpu_id") = -1);
+
+  m.def("get_cpu_arc_array1", []() -> Array1<Arc> {
+    std::vector<Arc> arcs = {
+        {1, 2, 3, 1.5},
+        {10, 20, 30, 2.5},
+    };
+    return Array1<Arc>(GetCpuContext(), arcs);
+  });
+
+  m.def(
+      "get_cuda_arc_array1",
+      [](int32_t gpu_id = -1) -> Array1<Arc> {
+        std::vector<Arc> arcs = {
+            {1, 2, 3, 1.5},
+            {10, 20, 30, 2.5},
+        };
+        return Array1<Arc>(GetCudaContext(gpu_id), arcs);
+      },
+      py::arg("gpu_id") = -1);
+
+  m.def("get_cpu_int_array2", []() -> Array2<int32_t> {
+    Array1<int32_t> array1(GetCpuContext(), {1, 2, 3, 4, 5, 6});
+    return Array2<int32_t>(array1, 2, 3);
+  });
+
+  m.def(
+      "get_cuda_float_array2",
+      [](int32_t gpu_id = -1) -> Array2<float> {
+        Array1<float> array1(GetCudaContext(gpu_id), {1, 2, 3, 4, 5, 6});
+        return Array2<float>(array1, 2, 3);
+      },
+      py::arg("gpu_id") = -1);
 }
 
 }  // namespace k2