Mark grad(sqrtm) as CPU only test

.github/workflows/tests.yml

@@ -57,7 +57,7 @@ jobs:
 
     - name: Run tests
       run: |
-        python3 -m pytest -m fast --memray -vv
+        python3 -m pytest -m "fast and not cpu" --memray -vv
 
   tests:
     name: Python ${{ matrix.python-version }} on ${{ matrix.os }}

pyproject.toml

@@ -113,8 +113,8 @@ testpaths = [
     "tests",
 ]
 markers = [
+    "cpu: Mark tests as CPU only.",
     "fast: Mark tests as fast.",
-    "notebook: Mark tests as notebook related.",
 ]
 
 [tool.coverage.run]

src/ott/problems/quadratic/quadratic_problem.py

@@ -301,7 +301,7 @@ def update_linearization(
     transport_matrix = transport.matrix * rescale_factor
 
     if not self.is_balanced:
-      # Rescales transport for Unbalanced GW according to Sejourne et al (2021).
+      # Rescales transport for Unbalanced GW according to Sejourne et al. (2021)
       transport_mass = jax.lax.stop_gradient(marginal_1.sum())
       epsilon = update_epsilon_unbalanced(epsilon, transport_mass)
       unbalanced_correction = self.cost_unbalanced_correction(

tests/math/matrix_square_root_test.py

@@ -171,6 +171,8 @@ def test_solve_bartels_stewart_batch(self):
     )
     np.testing.assert_allclose(self.x, x[0, 0], atol=1.e-5)
 
+  # gradient needs Schur decomposition, currently not implemented in jax for GPU
+  @pytest.mark.cpu
   @pytest.mark.fast.with_args(
       "fn,n_tests,dim,epsilon,atol,rtol",
       [(lambda x: matrix_square_root.sqrtm(x)[0], 3, 3, 1e-6, 1e-6, 1e-6),