Fix unused variables in tests, lexer and badge

README.md

@@ -1,4 +1,4 @@
-<img src="https://raw.githubusercontent.com/ott-jax/ott/main/docs/_static/logoOTT.png" width="10%" alt="logo">
+<img src="https://raw.githubusercontent.com/ott-jax/ott/main/docs/_static/images/logoOTT.png" width="10%" alt="logo">
 
 # Optimal Transport Tools (OTT).
 

docs/conf.py

@@ -93,8 +93,8 @@
 # a list of builtin themes.
 
 html_theme = 'sphinx_book_theme'
-html_logo = '_static/logoOTT.png'
-html_favicon = '_static/logoOTT.ico'
+html_logo = '_static/images/logoOTT.png'
+html_favicon = '_static/images/logoOTT.ico'
 
 # Add any paths that contain custom static files (such as style sheets) here,
 # relative to this directory. They are copied after the builtin static files,
@@ -110,7 +110,7 @@
     <div class="docutils container">
         <a class="reference external"
            href="https://colab.research.google.com/github/ott-jax/ott/blob/main/{{ docname|e }}">
-        <img alt="Colab badge" src="https://colab.research.google.com/assets/colab-badge.svg" width="125px">
+        <img alt="Open in Colab" src="../_static/images/colab-badge.svg" width="125px">
         </a>
     </div>
 """

docs/notebooks/One_Sinkhorn.ipynb

@@ -20,7 +20,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -32,7 +32,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
    "metadata": {
     "id": "CRSA3KsJtlpF"
    },
@@ -1134,7 +1134,16 @@
    "name": "python3"
   },
   "language_info": {
-   "name": "python"
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.12"
   }
  },
  "nbformat": 4,

docs/notebooks/fairness.ipynb

@@ -32,7 +32,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -592,7 +592,16 @@
    "name": "python3"
   },
   "language_info": {
-   "name": "python"
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.7"
   }
  },
  "nbformat": 4,

docs/notebooks/gromov_wasserstein.ipynb

@@ -22,22 +22,8 @@
   {
    "cell_type": "code",
    "execution_count": 1,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "yoauHUVPTyAl",
-    "outputId": "d8bc8507-972f-4ac9-e997-e733edc218be"
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Requirement already satisfied: ott-jax in /usr/local/lib/python3.7/dist-packages (0.2.2)\n"
-     ]
-    }
-   ],
+   "metadata": {},
+   "outputs": [],
    "source": [
     "import sys\n",
     "\n",
@@ -9154,12 +9140,21 @@
    "provenance": []
   },
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "ott",
    "language": "python",
-   "name": "python3"
+   "name": "ott"
   },
   "language_info": {
-   "name": "python"
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.12"
   }
  },
  "nbformat": 4,

docs/notebooks/soft_sort.ipynb

@@ -11,7 +11,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -23,7 +23,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
    "metadata": {
     "id": "CRSA3KsJtlpF"
    },
@@ -765,7 +765,16 @@
    "name": "python3"
   },
   "language_info": {
-   "name": "python"
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.9.12"
   }
  },
  "nbformat": 4,

tests/core/continuous_barycenter_test.py

@@ -201,7 +201,7 @@ def test_bures_barycenter(
     b1 = b2 = jnp.ones(num_components) / num_components
 
     y = jnp.concatenate((y1, y2))
-    jnp.concatenate((b1, b2))
+    b = jnp.concatenate((b1, b2))
 
     gmm_generator = gaussian_mixture.GaussianMixture.from_random(
         rng, n_components=bar_size, n_dimensions=dimension
@@ -213,7 +213,8 @@ def test_bures_barycenter(
     x_init = means_and_covs_to_x(x_init_means, x_init_covs, dimension)
 
     seg_y, seg_b = segment.segment_point_cloud(
-        y,
+        x=y,
+        a=b,
         num_segments=num_measures,
         max_measure_size=num_components,
         num_per_segment=(num_components, num_components),
@@ -279,6 +280,7 @@ def test_bures_barycenter_different_number_of_components(
     n_components = jnp.array([3, 4])  # the number of components of the GMMs
     num_measures = n_components.size
     bar_size = 5  # the size of the barycenter
+    max_measure_size = int(jnp.max(n_components))
 
     # Create an instance of the Bures cost class.
     b_cost = costs.Bures(dimension=dim)
@@ -345,7 +347,7 @@ def test_bures_barycenter_different_number_of_components(
         cost_fn=b_cost,
         epsilon=epsilon,
         num_segments=num_measures,
-        max_measure_size=4,
+        max_measure_size=max_measure_size,
         segment_ids=seg_ids,
     )
     assert bar_p.max_measure_size == 4

tests/tools/sinkhorn_divergence_test.py

@@ -261,7 +261,7 @@ def test_segment_sinkhorn_different_segment_sizes(self):
     segmented_divergences = sink_div(
         jnp.concatenate((x1, x2)),
         jnp.concatenate((y1, y2)),
-        # these 2 arguments are not necessary for jitting
+        # these 2 arguments are not necessary for jitting:
         # num_segments=2,
         # max_measure_size=15,
         num_per_segment_x=(10, 12),