merge

.flake8

@@ -0,0 +1,51 @@
+[flake8]
+max-line-length = 80
+ignore =
+    # line break before a binary operator -> black does not adhere to PEP8
+    W503
+    # line break occured after a binary operator -> black does not adhere to PEP8
+    W504
+    # line too long -> we accept long comment lines; black gets rid of long code lines
+    E501
+    # whitespace before : -> black does not adhere to PEP8
+    E203
+    # missing whitespace after ,', ';', or ':' -> black does not adhere to PEP8
+    E231
+    # continuation line over-indented for hanging indent -> black does not adhere to PEP8
+    E126
+    # E266 too many leading '#' for block comment -> this is fine for indicating sections
+    E262
+    # Do not assign a lambda expression, use a def -> lambda expression assignments are convenient
+    E731
+    # allow I, O, l as variable names -> I is the identity matrix, i, j, k, l is reasonable indexing notation
+    E741
+    # Missing docstring in public package
+    D104
+    # ... imported but unused
+    F401
+    # Missing docstring in public module
+    D100
+    # Missing docstring in __init__
+    D107
+    # Do not perform function calls in argument defaults.
+    B008
+    # line break before binary operator
+    W503
+    # Missing docstring in magic method
+    D105
+    # whitespace before ':'
+    E203
+    # format string does contain unindexed parameters
+    P101
+    # indentation is not a multiple of 4
+    E111, E114
+    # Missing blank line before section
+    D411
+exclude = .git,__pycache__,build,docs/_build,dist
+# C409: Unnecessary <dict/list/tuple> call - rewrite as a literal.
+per-file-ignores =
+    tests/*: D,C408
+    */__init__.py: F401
+    examples/*: D101, D102, D103
+    docs/*: D101, D102
+    src/ott/types.py: D102

.github/workflows/lint.yml

@@ -28,6 +28,13 @@ jobs:
       with:
         python-version: '3.10'
 
+    - name: Cache pre-commit
+      uses: actions/cache@v3
+      if: ${{ matrix.lint-kind == 'code' }}
+      with:
+        path: ~/.cache/pre-commit
+        key: pre-commit-${{ env.pythonLocation }}-${{ hashFiles('**/.pre-commit-config.yaml') }}
+
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
@@ -36,8 +43,8 @@ jobs:
     - name: Install PyEnchant
       if: ${{ matrix.lint-kind == 'docs' }}
       run: |
-        sudo apt update -y
-        sudo apt install libenchant-2-dev
+        sudo apt-get update -y
+        sudo apt-get install libenchant-2-dev
         python -m pip install pyenchant
 
     - name: Lint ${{ matrix.lint-kind }}

.github/workflows/tests.yml

@@ -23,6 +23,16 @@ jobs:
       with:
         python-version: '3.9'
 
+    - name: Install apt dependencies
+      if: ${{ runner.os == 'Linux' }}
+      run: |
+        sudo apt-get install libsuitesparse-dev
+
+    - name: Install brew dependencies
+      if: ${{ runner.os == 'macOS' }}
+      run: |
+        brew install suite-sparse
+
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
@@ -44,6 +54,7 @@ jobs:
       options: --gpus="device=12"
     steps:
     - uses: actions/checkout@v3
+
     - name: Install dependencies
       # `jax[cuda]<0.4` because of Docker issues: https://github.com/google/jax/issues/13758
       # `flax<0.6.5` because it requires `jax>=0.4.2`
@@ -80,6 +91,16 @@ jobs:
       with:
         python-version: ${{ matrix.python-version }}
 
+    - name: Install apt dependencies
+      if: ${{ runner.os == 'Linux' }}
+      run: |
+        sudo apt-get install libsuitesparse-dev
+
+    - name: Install brew dependencies
+      if: ${{ runner.os == 'macOS' }}
+      run: |
+        brew install suite-sparse
+
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip

docs/conf.py

@@ -1,10 +1,10 @@
-# Copyright 2022 Google LLC.
+# Copyright OTT-JAX
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
-#     http://www.apache.org/licenses/LICENSE-2.0
+#   http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,

docs/geometry.rst

@@ -52,6 +52,7 @@ Cost Functions
 .. autosummary::
     :toctree: _autosummary
 
+    costs.CostFn
     costs.SqPNorm
     costs.PNormP
     costs.SqEuclidean

docs/tools.rst

@@ -15,6 +15,13 @@ Segmented Sinkhorn
 
     segment_sinkhorn.segment_sinkhorn
 
+Plotting
+--------
+.. autosummary::
+    :toctree: _autosummary
+
+    plot.Plot
+
 Sinkhorn Divergence
 -------------------
 .. autosummary::

examples/fairness/config.py

@@ -1,10 +1,10 @@
-# Copyright 2022 Google LLC.
+# Copyright OTT-JAX
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
-#     http://www.apache.org/licenses/LICENSE-2.0
+#   http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,

examples/fairness/data.py

@@ -1,10 +1,10 @@
-# Copyright 2022 Google LLC.
+# Copyright OTT-JAX
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
-#     http://www.apache.org/licenses/LICENSE-2.0
+#   http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,

examples/fairness/losses.py

@@ -1,10 +1,10 @@
-# Copyright 2022 Google LLC.
+# Copyright OTT-JAX
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
-#     http://www.apache.org/licenses/LICENSE-2.0
+#   http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,

examples/fairness/main.py

@@ -1,10 +1,10 @@
-# Copyright 2022 Google LLC.
+# Copyright OTT-JAX
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
-#     http://www.apache.org/licenses/LICENSE-2.0
+#   http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,

examples/fairness/models.py

@@ -1,10 +1,10 @@
-# Copyright 2022 Google LLC.
+# Copyright OTT-JAX
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
-#     http://www.apache.org/licenses/LICENSE-2.0
+#   http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,

examples/fairness/train.py

@@ -1,10 +1,10 @@
-# Copyright 2022 Google LLC.
+# Copyright OTT-JAX
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
-#     http://www.apache.org/licenses/LICENSE-2.0
+#   http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,

examples/soft_error/config.py

@@ -1,10 +1,10 @@
-# Copyright 2022 Google LLC.
+# Copyright OTT-JAX
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
-#     http://www.apache.org/licenses/LICENSE-2.0
+#   http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,

examples/soft_error/data.py

@@ -1,10 +1,10 @@
-# Copyright 2022 Google LLC.
+# Copyright OTT-JAX
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
-#     http://www.apache.org/licenses/LICENSE-2.0
+#   http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,

examples/soft_error/losses.py

@@ -1,10 +1,10 @@
-# Copyright 2022 Google LLC.
+# Copyright OTT-JAX
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
-#     http://www.apache.org/licenses/LICENSE-2.0
+#   http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,

examples/soft_error/main.py

@@ -1,10 +1,10 @@
-# Copyright 2022 Google LLC.
+# Copyright OTT-JAX
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
-#     http://www.apache.org/licenses/LICENSE-2.0
+#   http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,

examples/soft_error/model.py

@@ -1,10 +1,10 @@
-# Copyright 2022 Google LLC.
+# Copyright OTT-JAX
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
-#     http://www.apache.org/licenses/LICENSE-2.0
+#   http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,

examples/soft_error/train.py

@@ -1,10 +1,10 @@
-# Copyright 2022 Google LLC.
+# Copyright OTT-JAX
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
-#     http://www.apache.org/licenses/LICENSE-2.0
+#   http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,

pyproject.toml

@@ -74,9 +74,6 @@ test = [
     "networkx>=2.5",
     "scikit-learn>=1.0"
 ]
-experimental = [
-    "scikit-sparse>=0.4.6",
-]
 docs = [
     "sphinx>=4.0",
     "ipython>=7.20.0",
@@ -177,6 +174,7 @@ legacy_tox_ini = """
     [testenv]
     extras = test
     pass_env = CUDA_*,PYTEST_*,CI
+    deps = py{38,39,310}: scikit-sparse>=0.4.6
     commands =
         gpu: pip install "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
         python -m pytest {tty:--color=yes} {posargs: \

src/ott/__init__.py

@@ -1,10 +1,10 @@
-# Copyright 2022 Google LLC.
+# Copyright OTT-JAX
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
-#     http://www.apache.org/licenses/LICENSE-2.0
+#   http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,

src/ott/geometry/__init__.py

@@ -1,10 +1,10 @@
-# Copyright 2022 Google LLC.
+# Copyright OTT-JAX
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
-#     http://www.apache.org/licenses/LICENSE-2.0
+#   http://www.apache.org/licenses/LICENSE-2.0
 #
 # Unless required by applicable law or agreed to in writing, software
 # distributed under the License is distributed on an "AS IS" BASIS,