Feature/fast tests

.codecov.yml

@@ -0,0 +1,18 @@
+codecov:
+  require_ci_to_pass: no
+  strict_yaml_branch: main
+
+coverage:
+  range: 75..100
+  status:
+    project:
+      default:
+        target: 1
+    patch: off
+
+comment:
+  layout: reach, diff, files
+  behavior: default
+  require_changes: true
+  branches:
+  - main

.github/workflows/tests.yml

@@ -14,24 +14,42 @@ jobs:
       matrix:
         python-version: ['3.8']
         os: [ubuntu-latest]
+        test_mark: [fast, all]
 
     steps:
     - uses: actions/checkout@v2
     - uses: actions/setup-python@v2
       with:
         python-version: ${{ matrix.python-version }}
+
     - name: Build
       run: |
         set -xe
         python -VV
-        pip install --upgrade pip setuptools wheel
+        pip install --upgrade pip setuptools
         pip install -e '.[test]'
-      shell: bash
-    - name: Run tests
+
+    - name: Print versions
       run: |
-        set -xe
         python -VV
         python -c "import jax; print('jax', jax.__version__)"
         python -c "import jaxlib; print('jaxlib', jaxlib.__version__)"
-        pytest tests -v --memray
-      shell: bash
+
+    - name: Run fast tests
+      if: ${{ matrix.test_mark == 'fast' }}
+      run: |
+        pytest -m fast --cov=ott --cov-report=xml --cov-report=term-missing --cov-config=setup.cfg --memray -v tests/
+
+    - name: Run all tests
+      if: ${{ matrix.test_mark == 'all' }}
+      run: |
+        pytest --cov=ott --cov-report=xml --cov-report=term-missing --cov-config=setup.cfg --memray -v tests/
+
+    - name: Upload coverage
+      uses: codecov/codecov-action@v3
+      with:
+        files: ./coverage.xml
+        flags: unittests-${{ matrix.test_mark }}
+        env_vars: OS,PYTHON
+        fail_ci_if_error: true
+        verbose: true

README.md

@@ -2,7 +2,8 @@
 
 # Optimal Transport Tools (OTT).
 
-![Tests](https://github.com/ott-jax/ott/actions/workflows/tests.yml/badge.svg)
+![Tests](https://img.shields.io/github/workflow/status/ott-jax/ott/tests/main)
+![Coverage](https://img.shields.io/codecov/c/github/ott-jax/ott/main)
 
 **See [full documentation](https://ott-jax.readthedocs.io/en/latest/).**
 

ott/tools/sinkhorn_divergence.py

@@ -13,19 +13,22 @@
 # limitations under the License.
 """Implements the sinkhorn divergence."""
 
-import collections
-from typing import Any, Dict, Mapping, Optional
+from typing import Any, Dict, List, Mapping, NamedTuple, Optional, Tuple
 
 import jax
 from jax import numpy as jnp
 
 from ott.core import segment, sinkhorn
 from ott.geometry import geometry, pointcloud
 
-SinkhornDivergenceOutput = collections.namedtuple(
-    'SinkhornDivergenceOutput',
-    ['divergence', 'potentials', 'geoms', 'errors', 'converged']
-)
+
+class SinkhornDivergenceOutput(NamedTuple):
+  divergence: float
+  potentials: Tuple[List[jnp.ndarray], List[jnp.ndarray], List[jnp.ndarray]]
+  geoms: Tuple[geometry.Geometry, geometry.Geometry, geometry.Geometry]
+  errors: Tuple[Optional[jnp.ndarray], Optional[jnp.ndarray],
+                Optional[jnp.ndarray]]
+  converged: Tuple[bool, bool, bool]
 
 
 def sinkhorn_divergence(
@@ -37,7 +40,7 @@ def sinkhorn_divergence(
     static_b: bool = False,
     share_epsilon: bool = True,
     **kwargs: Any,
-):
+) -> SinkhornDivergenceOutput:
   """Compute Sinkhorn divergence defined by a geometry, weights, parameters.
 
   Args:
@@ -82,10 +85,13 @@ def sinkhorn_divergence(
 
 
 def _sinkhorn_divergence(
-    geometry_xy: geometry.Geometry, geometry_xx: geometry.Geometry,
-    geometry_yy: Optional[geometry.Geometry], a: jnp.ndarray, b: jnp.ndarray,
-    **kwargs
-):
+    geometry_xy: geometry.Geometry,
+    geometry_xx: geometry.Geometry,
+    geometry_yy: Optional[geometry.Geometry],
+    a: jnp.ndarray,
+    b: jnp.ndarray,
+    **kwargs: Any,
+) -> SinkhornDivergenceOutput:
   """Compute the (unbalanced) sinkhorn divergence for the wrapper function.
 
     This definition includes a correction depending on the total masses of each
@@ -159,7 +165,7 @@ def segment_sinkhorn_divergence(
     sinkhorn_kwargs: Optional[Mapping[str, Any]] = None,
     static_b: bool = False,
     share_epsilon: bool = True,
-    **kwargs
+    **kwargs: Any
 ) -> jnp.ndarray:
   """Compute Sinkhorn divergence between subsets of data with point cloud.
 

setup.cfg

@@ -53,7 +53,8 @@ test =
     pytest
     pytest-xdist
     pytest-memray
-    absl-py
+    pytest-cov
+    chex
 docs =
     sphinx>=4.0
     nbsphinx>=0.8.0
@@ -63,3 +64,24 @@ docs =
     sphinx-book-theme
 dev =
     pre-commit
+
+[coverage:run]
+branch = true
+source = ott
+omit = */__init__.py
+
+[coverage:report]
+exclude_lines =
+    \#.*pragma:\s*no.?cover
+    ^if __name__ == .__main__.:$
+    ^\s*raise AssertionError\b
+    ^\s*raise NotImplementedError\b
+    ^\s*return NotImplemented\b
+precision = 2
+show_missing = True
+skip_empty = True
+sort = Miss
+
+[tool:pytest]
+markers =
+    fast: Mark tests as fast.

tests/conftest.py

@@ -0,0 +1,68 @@
+import collections.abc
+import itertools
+from typing import Any, Mapping, Optional, Sequence
+
+import jax
+import pytest
+from _pytest.python import Metafunc
+
+
+def pytest_generate_tests(metafunc: Metafunc) -> None:
+  if not hasattr(metafunc.function, "pytestmark"):
+    # no annotation
+    return
+
+  fast_marks = [m for m in metafunc.function.pytestmark if m.name == "fast"]
+  if fast_marks:
+    mark, = fast_marks
+    selected: Optional[Mapping[str, Any]] = mark.kwargs.pop("only_fast", None)
+    ids: Optional[Sequence[str]] = mark.kwargs.pop("ids", None)
+
+    if mark.args:
+      argnames, argvalues = mark.args
+    else:
+      argnames = tuple(mark.kwargs.keys())
+      argvalues = [
+          (vs,) if not isinstance(vs, (str, collections.abc.Iterable)) else vs
+          for vs in mark.kwargs.values()
+      ]
+      argvalues = list(itertools.product(*argvalues))
+
+    opt = str(metafunc.config.getoption("-m"))
+    if "fast" in opt:  # filter if `-m fast` was passed
+      if selected is None:
+        combinations = argvalues
+      elif isinstance(selected, dict):
+        combinations = []
+        for vs in argvalues:
+          if selected == dict(zip(argnames, vs)):
+            combinations.append(vs)
+      elif isinstance(selected, (tuple, list)):
+        # TODO(michalk8): support passing ids?
+        combinations = [argvalues[s] for s in selected]
+        ids = None if ids is None else [ids[s] for s in selected]
+      elif isinstance(selected, int):
+        combinations = [argvalues[selected]]
+        ids = None if ids is None else [ids[selected]]
+      else:
+        raise TypeError(f"Invalid fast selection type `{type(selected)}`.")
+    else:
+      combinations = argvalues
+
+    if argnames:
+      metafunc.parametrize(argnames, combinations, ids=ids)
+
+
+@pytest.fixture(scope="session")
+def rng():
+  return jax.random.PRNGKey(0)
+
+
+@pytest.fixture()
+def enable_x64():
+  previous_value = jax.config.jax_enable_x64
+  jax.config.update("jax_enable_x64", True)
+  try:
+    yield
+  finally:
+    jax.config.update("jax_enable_x64", previous_value)