Support printing ragged tensors in a more compact way. (#831)

* Support printing ragged tensors in a more compact way. * Disable support for torch 1.3.1 * Fix test failures.
k2-fsa · pkufool · Apr 14, 2022 · Sep 8, 2021 · Sep 14, 2021 · Sep 14, 2021
commit 8030001c9a002aa17e090a41de3f1146bdfe1e78
diff --git a/.github/workflows/build-doc.yml b/.github/workflows/build-doc.yml
@@ -23,6 +23,7 @@ on:
     branches:
     - master
     - doc
+    - doc-test
 
 env:
   # debug is faster in terms of compilation time

diff --git a/.github/workflows/nightly-cpu.yml b/.github/workflows/nightly-cpu.yml
@@ -40,13 +40,10 @@ jobs:
       matrix:
         os: [ubuntu-18.04, macos-10.15]
         # Python 3.9 is for PyTorch 1.7.1, 1.8.x, 1.9.0
-        # torch 1.3.1 supports only Python 3.5/6/7
         python-version: [3.6, 3.7, 3.8, 3.9]
-        torch: ["1.3.1", "1.4.0", "1.5.0", "1.5.1", "1.6.0", "1.7.0", "1.7.1", "1.8.0", "1.8.1", "1.9.0"]
+        torch: ["1.4.0", "1.5.0", "1.5.1", "1.6.0", "1.7.0", "1.7.1", "1.8.0", "1.8.1", "1.9.0"]
         exclude:
-          - python-version: 3.9 # exclude Python 3.9 for [1.3.1, 1.4.0, 1.5.0, 1.5.1, 1.6.0, 1.7.0]
-            torch: "1.3.1"
-          - python-version: 3.9
+          - python-version: 3.9 # exclude Python 3.9 for [1.4.0, 1.5.0, 1.5.1, 1.6.0, 1.7.0]
             torch: "1.4.0"
           - python-version: 3.9
             torch: "1.5.0"
@@ -56,8 +53,6 @@ jobs:
             torch: "1.6.0"
           - python-version: 3.9
             torch: "1.7.0"
-          - python-version: 3.8 # exclude Python 3.8 for [1.3.1]
-            torch: "1.3.1"
 
     steps:
       - uses: actions/checkout@v2

diff --git a/docs/source/conf.py b/docs/source/conf.py
@@ -14,15 +14,14 @@
 import re
 import sys
 sys.path.insert(0, os.path.abspath('../../k2/python'))
-sys.path.insert(0, os.path.abspath('../../build-ragged/lib'))
 sys.path.insert(0, os.path.abspath('../../build/lib'))
 
 import sphinx_rtd_theme
 
 # -- Project information -----------------------------------------------------
 
 project = 'k2'
-copyright = '2020, k2 development team'
+copyright = '2020-2021, k2 development team'
 author = 'k2 development team'
 
 

diff --git a/docs/source/installation/for_developers.rst b/docs/source/installation/for_developers.rst
@@ -9,7 +9,7 @@ First, you have to install CMake, CUDA toolkit (with cuDNN), and PyTorch:
   - CMake 3.11.0 and 3.18.0 are known to work. Other CMake versions may work
     but they are not tested.
 
-  - Install PyTorch. PyTorch 1.5.x and above are known to work. Other PyTorch
+  - Install PyTorch. PyTorch 1.4.x and above are known to work. Other PyTorch
     versions may work, but they are not tested.
 
   - Install CUDA toolkit and cuDNN. CUDA 10.1 and above are known to work.
@@ -43,7 +43,7 @@ To build a release version, use:
   python3 -c "import k2; print(k2.__file__)"
   # It should print /some/path/k2/k2/python/k2/__init.py
 
-  python3 -c "import _k2; print(_k2.__file__)"
+  python3 -c "import torch; import _k2; print(_k2.__file__)"
   # It should print /some/path/k2/build_release/lib/_k2.cpython-38-x86_64-linux-gnu.so
   # (I assume that you're using Python 3.8, so there is a string 38 above)
 
@@ -63,10 +63,45 @@ To build a debug version, use:
   python3 -c "import k2; print(k2.__file__)"
   # It should print /some/path/k2/k2/python/k2/__init.py
 
-  python3 -c "import _k2; print(_k2.__file__)"
+  python3 -c "import torch; import _k2; print(_k2.__file__)"
   # It should print /some/path/k2/build_debug/lib/_k2.cpython-38-x86_64-linux-gnu.so
   # (I assume that you're using Python 3.8, so there is a string 38 above)
 
+.. HINT::
+
+  You can pass the option ``-DK2_WITH_CUDA=OFF`` to ``cmake`` to build
+  a CPU only version of k2.
+
+  It is much faster to build a CPU version than that of building a CUDA
+  version. When you are adding new features to k2, we recommend you to
+  create a diretory to build a CPU version to test your code. Once it is
+  working on CPU, you can create a new directory to build a CUDA version
+  to test your code.
+
+  That is, while adding and testing new features, use:
+
+    .. code-block:: bash
+
+      cd k2
+      mkdir build-cpu
+      cd build-cpu
+      cmake -DK2_WITH_CUDA=OFF -DCMAKE_BUILD_TYPE=Debug ..
+      make -j5
+      export PYTHONPATH=$PWD/../k2/python:$PWD/lib:$PYTHONPATH
+      # make test # to test your code
+
+  After it is working for CPU, you can use:
+
+    .. code-block:: bash
+
+      cd k2
+      mkdir build-cuda
+      cd build-cuda
+      cmake -DCMAKE_BUILD_TYPE=Debug ..
+      make -j5
+      export PYTHONPATH=$PWD/../k2/python:$PWD/lib:$PYTHONPATH
+      # make test # to test your code
+
 To run tests, use:
 
 .. code-block:: bash
@@ -154,16 +189,16 @@ To run a specific Python test, use:
 
   To check whether you are using a release version or a debug version, run:
 
-    .. code-block::
+    .. code-block:: bash
 
-      python3 -c "import _k2; print(_k2.__file__)"
+      python3 -c "import torch; import _k2; print(_k2.__file__)"
 
   It should print the directory where k2 was built. That is,
   the above output contains a string ``build_release`` or ``build_debug``.
 
   Alternatively, you can run:
 
-    .. code-block::
+    .. code-block:: bash
 
       python3 -m k2.version
 

diff --git a/docs/source/installation/pip.rst b/docs/source/installation/pip.rst
@@ -29,77 +29,80 @@ versions of Python, CUDA, and PyTorch.
   automagically. You don't need to pre-install PyTorch and cudatoolkit when using
   ``conda install``.
 
-The following commands install k2 with different CUDA versions:
+The following commands install k2 with different versions of CUDA and PyTorch:
 
 .. code-block:: bash
 
-  # Install k2 0.3.3 with CUDA 10.2 built on 20210509
+  # Install k2 1.8 with CUDA 10.1 built on 20210916
   #
-  # cu102 means CUDA 10.2
+  # You don't need to specifiy the Python version
   #
-  pip install k2==0.3.3+cu102.dev20210509 -f https://k2-fsa.org/nightly/
+  pip install k2==1.8.dev20210916+cuda10.1.torch1.7.1 -f https://k2-fsa.org/nightly/
 
-  # Install k2 0.3.3 with CUDA 11.0 built on 20210509
+  # Install k2 1.8 with CUDA 10.2 built on 20210916
   #
-  # cu110 means CUDA 11.0
   #
-  pip install k2==0.3.3+cu110.dev20210509 -f https://k2-fsa.org/nightly/
+  pip install k2==1.8.dev20210916+cuda10.2.torch1.7.1 -f https://k2-fsa.org/nightly/
 
-  # Install k2 0.3.3 with CUDA 10.1 built on 20210509
+  # Install k2 1.8 with CUDA 11.0 built on 20210916
   #
-  # CAUTION: you don't need to specify cu101 since CUDA 10.1 is the default
-  # CUDA version
-  #
-  pip install k2==0.3.3.dev20210509 -f https://k2-fsa.org/nightly/
+  pip install k2==1.8.dev20210916+cuda11.0.torch1.7.1 -f https://k2-fsa.org/nightly/
 
-  #
-  # dev20210509 means that version is built on 2021.05.09
-  #
   # Please always select the latest version. That is, the version
   # with the latest date.
 
+.. Caution::
+
+  We only provide pre-compiled versions of k2 with torch 1.7.1. If you need
+  other versions of PyTorch, please consider one of the following alternatives
+  to install k2:
+
+    - :ref:`install using conda`
+    - :ref:`install k2 from source`
+
 The following is the log for installing k2:
 
 .. code-block::
 
-  $ pip install k2==0.3.3.dev20210509 -f https://k2-fsa.org/nightly/
-  Looking in links: https://k2-fsa.org/nightly/
-  Collecting k2==0.3.3.dev20210509
-    Downloading https://k2-fsa.org/nightly/whl/k2-0.3.3.dev20210509-cp38-cp38-linux_x86_64.whl (54.4 MB)
-       |________________________________| 54.4 MB 487 kB/s
-  Requirement already satisfied: torch in ./py38/lib/python3.8/site-packages (from k2==0.3.3.dev20210509) (1.7.1+cu101)
-  Requirement already satisfied: graphviz in ./py38/lib/python3.8/site-packages (from k2==0.3.3.dev20210509) (0.15)
-  Requirement already satisfied: numpy in ./py38/lib/python3.8/site-packages (from torch->k2==0.3.3.dev20210509) (1.19.5)
-  Requirement already satisfied: typing-extensions in ./py38/lib/python3.8/site-packages (from torch->k2==0.3.3.dev20210509) (3.7.4.3)
-  Installing collected packages: k2
-  Successfully installed k2-0.3.3.dev20210509
-  WARNING: You are using pip version 21.0.1; however, version 21.1.1 is available.
-  You should consider upgrading via the '/xxx/bin/python3.8 -m pip install --upgrade pip' command.
+  $ pip install k2==1.8.dev20210916+cuda10.1.torch1.7.1 -f https://k2-fsa.org/nightly
+
+  Looking in links: https://k2-fsa.org/nightly
+  Collecting k2==1.8.dev20210916+cuda10.1.torch1.7.1
+    Downloading https://k2-fsa.org/nightly/whl/k2-1.8.dev20210916%2Bcuda10.1.torch1.7.1-cp38-cp38-linux_x86_64.whl (77.7 MB)
+       |________________________________| 77.7 MB 1.6 MB/s
+  Collecting torch==1.7.1
+    Using cached torch-1.7.1-cp38-cp38-manylinux1_x86_64.whl (776.8 MB)
+  Collecting graphviz
+    Using cached graphviz-0.17-py3-none-any.whl (18 kB)
+  Collecting typing-extensions
+    Downloading typing_extensions-3.10.0.2-py3-none-any.whl (26 kB)
+  Collecting numpy
+    Using cached numpy-1.21.2-cp38-cp38-manylinux_2_12_x86_64.manylinux2010_x86_64.whl (15.8 MB)
+  Installing collected packages: typing-extensions, numpy, torch, graphviz, k2
+  Successfully installed graphviz-0.17 k2-1.8.dev20210916+cuda10.1.torch1.7.1 numpy-1.21.2 torch-1.7.1 typing-extensions-3.10.0.2
 
 To verify that k2 is installed successfully, run:
 
 .. code-block::
 
   $ python3 -m k2.version
-  /xxx/lib/python3.8/runpy.py:127: RuntimeWarning: 'k2.version' found in sys.modules after import of package 'k2', but prior to execution of 'k2.version'; this may result in unpredictable behaviour
-    warn(RuntimeWarning(msg))
-  Collecting environment information...
 
-  k2 version: 0.3.3
+  k2 version: 1.8
   Build type: Release
-  Git SHA1: 8e2fa82dca767782351fec57ec187aa04015dcf2
-  Git date: Thu May 6 18:55:15 2021
+  Git SHA1: 646704e142438bcd1aaf4a6e32d95e5ccd93a174
+  Git date: Thu Sep 16 13:05:12 2021
   Cuda used to build k2: 10.1
   cuDNN used to build k2: 8.0.2
   Python version used to build k2: 3.8
   OS used to build k2: Ubuntu 18.04.5 LTS
-  CMake version: 3.20.2
+  CMake version: 3.21.2
   GCC version: 7.5.0
-  CMAKE_CUDA_FLAGS:  -D_GLIBCXX_USE_CXX11_ABI=0  --expt-extended-lambda -gencode arch=compute_35,code=sm_35 --expt-extended-lambda -gencode arch=compute_50,code=sm_50 --expt-extended-lambda -gencode arch=compute_60,code=sm_60 --expt-extended-lambda -gencode arch=compute_61,code=sm_61 --expt-extended-lambda -gencode arch=compute_70,code=sm_70 --expt-extended-lambda -gencode arch=compute_75,code=sm_75 --compiler-options -Wall --compiler-options -Wno-unknown-pragmas
-  CMAKE_CXX_FLAGS:  -D_GLIBCXX_USE_CXX11_ABI=0
+  CMAKE_CUDA_FLAGS:  --expt-extended-lambda -gencode arch=compute_35,code=sm_35 --expt-extended-lambda -gencode arch=compute_50,code=sm_50 --expt-extended-lambda -gencode arch=compute_60,code=sm_60 --expt-extended-lambda -gencode arch=compute_61,code=sm_61 --expt-extended-lambda -gencode arch=compute_70,code=sm_70 --expt-extended-lambda -gencode arch=compute_75,code=sm_75 -D_GLIBCXX_USE_CXX11_ABI=0 --compiler-options -Wall --compiler-options -Wno-unknown-pragmas --compiler-options -Wno-strict-overflow
+  CMAKE_CXX_FLAGS:  -D_GLIBCXX_USE_CXX11_ABI=0 -Wno-strict-overflow
   PyTorch version used to build k2: 1.7.1+cu101
   PyTorch is using Cuda: 10.1
   NVTX enabled: True
+  With CUDA: True
   Disable debug: True
   Sync kernels : False
   Disable checks: False

diff --git a/docs/source/installation/pip_pypi.rst b/docs/source/installation/pip_pypi.rst
@@ -31,45 +31,40 @@ The following command installs k2 from PyPI:
 
 .. code-block:: bash
 
-  pip install --pre k2
+  pip install k2
 
-The wheel packages on PyPI are built using `torch==1.7.1+cu101` on Ubuntu 18.04.
-If you are using other Linux systems or a different PyTorch version,
-the pre-built wheel packages may NOT work on your system, please install
-k2 from source in this case.
+.. Caution::
 
-.. CAUTION::
+  The wheel packages on PyPI are built using `torch==1.7.1+cu101` on Ubuntu 18.04.
+  If you are using other Linux systems or a different PyTorch version, the
+  pre-built wheel packages may NOT work on your system, please consider one of
+  the following alternatives to install k2:
 
-    k2 is still under active development and we are trying to keep
-    the packages on PyPI up to date. Please use ``--pre`` in ``pip install``.
-
-    If you want to try the latest version, please refer to
-    :ref:`install k2 from source`.
+      - :ref:`install using conda`
+      - :ref:`install k2 from source`
 
 To verify that k2 is installed successfully, run:
 
 .. code-block::
 
   $ python3 -m k2.version
-  /xxx/lib/python3.8/runpy.py:127: RuntimeWarning: 'k2.version' found in sys.modules after import of package 'k2', but prior to execution of 'k2.version'; this may result in unpredictable behaviour
-    warn(RuntimeWarning(msg))
-  Collecting environment information...
 
-  k2 version: 0.3.3
+  k2 version: 1.8
   Build type: Release
-  Git SHA1: d66cad5067563bb87710a40cf401af35cae816ff
-  Git date: Fri Apr 30 13:33:47 2021
+  Git SHA1: 646704e142438bcd1aaf4a6e32d95e5ccd93a174
+  Git date: Thu Sep 16 13:05:12 2021
   Cuda used to build k2: 10.1
   cuDNN used to build k2: 8.0.2
   Python version used to build k2: 3.8
   OS used to build k2: Ubuntu 18.04.5 LTS
-  CMake version: 3.20.1
-  GCC version: 5.5.0
-  CMAKE_CUDA_FLAGS:  -D_GLIBCXX_USE_CXX11_ABI=0  --expt-extended-lambda -gencode arch=compute_35,code=sm_35 --expt-extended-lambda -gencode arch=compute_50,code=sm_50 --expt-extended-lambda -gencode arch=compute_60,code=sm_60 --expt-extended-lambda -gencode arch=compute_61,code=sm_61 --expt-extended-lambda -gencode arch=compute_70,code=sm_70 --expt-extended-lambda -gencode arch=compute_75,code=sm_75 --compiler-options -Wall --compiler-options -Wno-unknown-pragmas
-  CMAKE_CXX_FLAGS:  -D_GLIBCXX_USE_CXX11_ABI=0
+  CMake version: 3.21.2
+  GCC version: 7.5.0
+  CMAKE_CUDA_FLAGS:  --expt-extended-lambda -gencode arch=compute_35,code=sm_35 --expt-extended-lambda -gencode arch=compute_50,code=sm_50 --expt-extended-lambda -gencode arch=compute_60,code=sm_60 --expt-extended-lambda -gencode arch=compute_61,code=sm_61 --expt-extended-lambda -gencode arch=compute_70,code=sm_70 --expt-extended-lambda -gencode arch=compute_75,code=sm_75 -D_GLIBCXX_USE_CXX11_ABI=0 --compiler-options -Wall --compiler-options -Wno-unknown-pragmas --compiler-options -Wno-strict-overflow
+  CMAKE_CXX_FLAGS:  -D_GLIBCXX_USE_CXX11_ABI=0 -Wno-strict-overflow
   PyTorch version used to build k2: 1.7.1+cu101
   PyTorch is using Cuda: 10.1
   NVTX enabled: True
+  With CUDA: True
   Disable debug: True
   Sync kernels : False
   Disable checks: False

diff --git a/k2/python/csrc/torch/v2/any.cu b/k2/python/csrc/torch/v2/any.cu
@@ -70,6 +70,13 @@ void PybindRaggedAny(py::module &m) {
       [](const RaggedAny &self) -> std::string { return self.ToString(); },
       kRaggedAnyStrDoc);
 
+  any.def(
+      "to_str_simple",
+      [](const RaggedAny &self) -> std::string {
+        return self.ToString(/*compact*/ true);
+      },
+      kRaggedAnyToStrSimpleDoc);
+
   any.def(
       "__repr__",
       [](const RaggedAny &self) -> std::string { return self.ToString(); },

diff --git a/k2/python/csrc/torch/v2/doc/any.h b/k2/python/csrc/torch/v2/doc/any.h
@@ -595,6 +595,28 @@ RaggedTensor([[1],
 RaggedTensor([[1, 2]], device='cuda:0', dtype=torch.int32)
 )doc";
 
+static constexpr const char *kRaggedAnyToStrSimpleDoc = R"doc(
+Convert a ragged tensor to a string representation, which
+is more compact than ``self.__str__``.
+
+An example output is given below::
+
+  RaggedTensor([[[1, 2, 3], [], [0]], [[2], [3, 10.5]]], dtype=torch.float32)
+
+>>> import k2.ragged as k2r
+>>> a = k2r.RaggedTensor([ [[1, 2, 3], [], [0]], [[2], [3, 10.5]] ])
+>>> a
+RaggedTensor([[[1, 2, 3],
+               [],
+               [0]],
+              [[2],
+               [3, 10.5]]], dtype=torch.float32)
+>>> str(a)
+'RaggedTensor([[[1, 2, 3],\n               [],\n               [0]],\n              [[2],\n               [3, 10.5]]], dtype=torch.float32)'
+>>> a.to_str_simple()
+'RaggedTensor([[[1, 2, 3], [], [0]], [[2], [3, 10.5]]], dtype=torch.float32)'
+)doc";
+
 static constexpr const char *kRaggedAnyGetItemDoc = R"doc(
 Select the i-th sublist along axis 0.