From f840639ca441de93f5da3b9700c5bec813a5684d Mon Sep 17 00:00:00 2001
From: Michal Klein <46717574+michalk8@users.noreply.github.com>
Date: Mon, 27 Feb 2023 20:00:27 +0100
Subject: [PATCH] Remove module docs, add missing copyright

---
 src/ott/_version.py                                 | 13 +++++++++++++
 src/ott/problems/linear/barycenter_problem.py       |  1 -
 src/ott/problems/linear/linear_problem.py           |  2 --
 src/ott/solvers/linear/discrete_barycenter.py       |  2 --
 src/ott/tools/segment_sinkhorn.py                   |  1 -
 src/ott/tools/sinkhorn_divergence.py                |  1 -
 src/ott/tools/soft_sort.py                          |  2 --
 src/ott/types.py                                    |  2 ++
 tests/conftest.py                                   | 13 +++++++++++++
 tests/geometry/costs_test.py                        |  1 -
 tests/geometry/graph_test.py                        | 13 +++++++++++++
 tests/geometry/pointcloud_test.py                   |  1 -
 tests/geometry/scaling_cost_test.py                 |  1 -
 tests/geometry/subsetting_test.py                   | 13 +++++++++++++
 tests/initializers/linear/sinkhorn_init_test.py     |  1 -
 tests/initializers/linear/sinkhorn_lr_init_test.py  |  2 --
 tests/initializers/quadratic/gw_init_test.py        |  2 --
 tests/math/lse_test.py                              |  2 --
 tests/math/matrix_square_root_test.py               |  1 -
 tests/problems/linear/potentials_test.py            | 13 +++++++++++++
 tests/solvers/linear/continuous_barycenter_test.py  |  1 -
 tests/solvers/linear/discrete_barycenter_test.py    |  1 -
 tests/solvers/linear/sinkhorn_diff_test.py          |  1 -
 tests/solvers/linear/sinkhorn_grid_test.py          |  2 --
 tests/solvers/linear/sinkhorn_lr_test.py            |  1 -
 tests/solvers/linear/sinkhorn_misc_test.py          |  1 -
 tests/solvers/linear/sinkhorn_test.py               |  1 -
 tests/solvers/nn/icnn_test.py                       |  2 --
 tests/solvers/nn/neuraldual_test.py                 |  2 --
 tests/solvers/quadratic/fgw_test.py                 |  1 -
 tests/solvers/quadratic/gw_barycenter_test.py       |  1 -
 tests/solvers/quadratic/gw_test.py                  |  1 -
 tests/tools/gaussian_mixture/fit_gmm_pair_test.py   |  2 --
 tests/tools/gaussian_mixture/fit_gmm_test.py        |  2 --
 .../gaussian_mixture/gaussian_mixture_pair_test.py  |  2 --
 .../tools/gaussian_mixture/gaussian_mixture_test.py |  2 --
 tests/tools/gaussian_mixture/gaussian_test.py       |  2 --
 tests/tools/gaussian_mixture/linalg_test.py         |  2 --
 tests/tools/gaussian_mixture/probabilities_test.py  |  2 --
 tests/tools/gaussian_mixture/scale_tril_test.py     |  2 --
 tests/tools/k_means_test.py                         | 13 +++++++++++++
 tests/tools/segment_sinkhorn_test.py                |  2 --
 tests/tools/sinkhorn_divergence_test.py             |  1 -
 tests/tools/soft_sort_test.py                       |  1 -
 tests/utils_test.py                                 | 13 +++++++++++++
 45 files changed, 93 insertions(+), 55 deletions(-)

diff --git a/src/ott/_version.py b/src/ott/_version.py
index 85ed093e4..972f66cbd 100644
--- a/src/ott/_version.py
+++ b/src/ott/_version.py
@@ -1,3 +1,16 @@
+# 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
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 from importlib.metadata import PackageNotFoundError, version
 
 try:
diff --git a/src/ott/problems/linear/barycenter_problem.py b/src/ott/problems/linear/barycenter_problem.py
index 65b612b5a..b7ed68082 100644
--- a/src/ott/problems/linear/barycenter_problem.py
+++ b/src/ott/problems/linear/barycenter_problem.py
@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Classes defining OT problem(s) (objective function + utilities)."""
 from typing import Any, Dict, Optional, Sequence, Tuple
 
 import jax
diff --git a/src/ott/problems/linear/linear_problem.py b/src/ott/problems/linear/linear_problem.py
index 828fd6740..f5f87c594 100644
--- a/src/ott/problems/linear/linear_problem.py
+++ b/src/ott/problems/linear/linear_problem.py
@@ -11,8 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Classes defining OT problem(s) (objective function + utilities)."""
-
 from typing import Any, Callable, Dict, Optional, Sequence, Tuple
 
 import jax
diff --git a/src/ott/solvers/linear/discrete_barycenter.py b/src/ott/solvers/linear/discrete_barycenter.py
index 5a8248801..dcfdc1470 100644
--- a/src/ott/solvers/linear/discrete_barycenter.py
+++ b/src/ott/solvers/linear/discrete_barycenter.py
@@ -11,8 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Implementation of :cite:`janati:20` Wasserstein barycenter algorithm."""
-
 import functools
 from typing import NamedTuple, Optional, Sequence
 
diff --git a/src/ott/tools/segment_sinkhorn.py b/src/ott/tools/segment_sinkhorn.py
index d998850f9..852afd24d 100644
--- a/src/ott/tools/segment_sinkhorn.py
+++ b/src/ott/tools/segment_sinkhorn.py
@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Segmented sinkhorn utility."""
 from types import MappingProxyType
 from typing import Any, Mapping, Optional, Tuple
 
diff --git a/src/ott/tools/sinkhorn_divergence.py b/src/ott/tools/sinkhorn_divergence.py
index e5cd8885f..54bfcb3ef 100644
--- a/src/ott/tools/sinkhorn_divergence.py
+++ b/src/ott/tools/sinkhorn_divergence.py
@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Implements the sinkhorn divergence."""
 from types import MappingProxyType
 from typing import Any, List, Mapping, NamedTuple, Optional, Tuple, Type
 
diff --git a/src/ott/tools/soft_sort.py b/src/ott/tools/soft_sort.py
index a15c04ac1..226d44bf8 100644
--- a/src/ott/tools/soft_sort.py
+++ b/src/ott/tools/soft_sort.py
@@ -11,8 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Soft sort operators."""
-
 import functools
 from typing import Any, Callable, Optional
 
diff --git a/src/ott/types.py b/src/ott/types.py
index 10210690d..7a4c88716 100644
--- a/src/ott/types.py
+++ b/src/ott/types.py
@@ -15,6 +15,8 @@
 
 import jax.numpy as jnp
 
+__all__ = ["Transport"]
+
 # TODO(michalk8): introduce additional types here
 
 
diff --git a/tests/conftest.py b/tests/conftest.py
index 9bb8eaf09..c2974d6ea 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,3 +1,16 @@
+# 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
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 import collections.abc
 import itertools
 from typing import Any, Mapping, Optional, Sequence
diff --git a/tests/geometry/costs_test.py b/tests/geometry/costs_test.py
index dc8581ee3..80525ed63 100644
--- a/tests/geometry/costs_test.py
+++ b/tests/geometry/costs_test.py
@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for the cost/norm functions."""
 from typing import Type
 
 import jax
diff --git a/tests/geometry/graph_test.py b/tests/geometry/graph_test.py
index d5d27119e..057120431 100644
--- a/tests/geometry/graph_test.py
+++ b/tests/geometry/graph_test.py
@@ -1,3 +1,16 @@
+# 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
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 import time
 from typing import Any, Callable, Literal, Optional, Tuple, Union
 
diff --git a/tests/geometry/pointcloud_test.py b/tests/geometry/pointcloud_test.py
index d93ae6129..a1048c448 100644
--- a/tests/geometry/pointcloud_test.py
+++ b/tests/geometry/pointcloud_test.py
@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for apply_cost and apply_kernel."""
 from typing import Union
 
 import jax
diff --git a/tests/geometry/scaling_cost_test.py b/tests/geometry/scaling_cost_test.py
index 652a3b5c8..2652dd905 100644
--- a/tests/geometry/scaling_cost_test.py
+++ b/tests/geometry/scaling_cost_test.py
@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for the option to scale the cost matrix."""
 from typing import Optional, Union
 
 import jax
diff --git a/tests/geometry/subsetting_test.py b/tests/geometry/subsetting_test.py
index f3fba9188..9edb04246 100644
--- a/tests/geometry/subsetting_test.py
+++ b/tests/geometry/subsetting_test.py
@@ -1,3 +1,16 @@
+# 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
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 from typing import Optional, Sequence, Tuple, Type, Union
 
 import jax
diff --git a/tests/initializers/linear/sinkhorn_init_test.py b/tests/initializers/linear/sinkhorn_init_test.py
index e2c00f8c7..9c996e213 100644
--- a/tests/initializers/linear/sinkhorn_init_test.py
+++ b/tests/initializers/linear/sinkhorn_init_test.py
@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for Sinkhorn initializers."""
 from typing import Literal, Optional
 
 import jax
diff --git a/tests/initializers/linear/sinkhorn_lr_init_test.py b/tests/initializers/linear/sinkhorn_lr_init_test.py
index 318d40c02..63ce10fc3 100644
--- a/tests/initializers/linear/sinkhorn_lr_init_test.py
+++ b/tests/initializers/linear/sinkhorn_lr_init_test.py
@@ -11,8 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for Sinkhorn initializers."""
-
 import jax
 import jax.numpy as jnp
 import numpy as np
diff --git a/tests/initializers/quadratic/gw_init_test.py b/tests/initializers/quadratic/gw_init_test.py
index 5dcf20024..b9aef30d8 100644
--- a/tests/initializers/quadratic/gw_init_test.py
+++ b/tests/initializers/quadratic/gw_init_test.py
@@ -11,8 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for Gromov-Wasserstein initializers."""
-
 import jax
 import numpy as np
 import pytest
diff --git a/tests/math/lse_test.py b/tests/math/lse_test.py
index b73f710a3..9a01d54cd 100644
--- a/tests/math/lse_test.py
+++ b/tests/math/lse_test.py
@@ -11,8 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for the jvp of a custom implementation of lse."""
-
 import jax
 import jax.numpy as jnp
 import numpy as np
diff --git a/tests/math/matrix_square_root_test.py b/tests/math/matrix_square_root_test.py
index 40bf0d6e3..fd4584c18 100644
--- a/tests/math/matrix_square_root_test.py
+++ b/tests/math/matrix_square_root_test.py
@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for matrix square roots."""
 from typing import Any, Callable
 
 import jax
diff --git a/tests/problems/linear/potentials_test.py b/tests/problems/linear/potentials_test.py
index 8be38168a..a67876287 100644
--- a/tests/problems/linear/potentials_test.py
+++ b/tests/problems/linear/potentials_test.py
@@ -1,3 +1,16 @@
+# 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
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 import jax
 import jax.numpy as jnp
 import numpy as np
diff --git a/tests/solvers/linear/continuous_barycenter_test.py b/tests/solvers/linear/continuous_barycenter_test.py
index b1b9eb913..aaf243cc8 100644
--- a/tests/solvers/linear/continuous_barycenter_test.py
+++ b/tests/solvers/linear/continuous_barycenter_test.py
@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for continuous barycenter."""
 import functools
 from typing import Tuple
 
diff --git a/tests/solvers/linear/discrete_barycenter_test.py b/tests/solvers/linear/discrete_barycenter_test.py
index a97d1c549..0bb94bc6d 100644
--- a/tests/solvers/linear/discrete_barycenter_test.py
+++ b/tests/solvers/linear/discrete_barycenter_test.py
@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
 import jax.numpy as jnp
 import pytest
 from ott.geometry import grid, pointcloud
diff --git a/tests/solvers/linear/sinkhorn_diff_test.py b/tests/solvers/linear/sinkhorn_diff_test.py
index 5c739f4e7..8512dcdae 100644
--- a/tests/solvers/linear/sinkhorn_diff_test.py
+++ b/tests/solvers/linear/sinkhorn_diff_test.py
@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for the differentiability of reg_ot_cost w.r.t weights/locations."""
 import functools
 from typing import Callable, List, Optional, Tuple
 
diff --git a/tests/solvers/linear/sinkhorn_grid_test.py b/tests/solvers/linear/sinkhorn_grid_test.py
index ba7474e74..abeeb9f88 100644
--- a/tests/solvers/linear/sinkhorn_grid_test.py
+++ b/tests/solvers/linear/sinkhorn_grid_test.py
@@ -11,8 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for Sinkhorn when applied on a grid."""
-
 import jax
 import jax.numpy as jnp
 import numpy as np
diff --git a/tests/solvers/linear/sinkhorn_lr_test.py b/tests/solvers/linear/sinkhorn_lr_test.py
index aea6a8a36..7cda2bb62 100644
--- a/tests/solvers/linear/sinkhorn_lr_test.py
+++ b/tests/solvers/linear/sinkhorn_lr_test.py
@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests Sinkhorn Low-Rank solver with various initializations."""
 import jax
 import jax.numpy as jnp
 import numpy as np
diff --git a/tests/solvers/linear/sinkhorn_misc_test.py b/tests/solvers/linear/sinkhorn_misc_test.py
index 250ac3e43..2787ca73f 100644
--- a/tests/solvers/linear/sinkhorn_misc_test.py
+++ b/tests/solvers/linear/sinkhorn_misc_test.py
@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests Anderson acceleration for Sinkhorn."""
 from typing import Optional, Tuple
 
 import chex
diff --git a/tests/solvers/linear/sinkhorn_test.py b/tests/solvers/linear/sinkhorn_test.py
index 026913624..bf0ff23d7 100644
--- a/tests/solvers/linear/sinkhorn_test.py
+++ b/tests/solvers/linear/sinkhorn_test.py
@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for Sinkhorn."""
 from typing import Any, Optional, Tuple
 
 import jax
diff --git a/tests/solvers/nn/icnn_test.py b/tests/solvers/nn/icnn_test.py
index 848c15b65..131d17b74 100644
--- a/tests/solvers/nn/icnn_test.py
+++ b/tests/solvers/nn/icnn_test.py
@@ -11,8 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for ICNN network architecture."""
-
 import jax
 import jax.numpy as jnp
 import numpy as np
diff --git a/tests/solvers/nn/neuraldual_test.py b/tests/solvers/nn/neuraldual_test.py
index 1dde9a69b..bec9db919 100644
--- a/tests/solvers/nn/neuraldual_test.py
+++ b/tests/solvers/nn/neuraldual_test.py
@@ -1,4 +1,3 @@
-#
 # Copyright OTT-JAX
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
@@ -12,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for implementation of ICNN-based Kantorovich dual by Makkuva+(2020)."""
 from typing import Sequence, Tuple
 
 import jax
diff --git a/tests/solvers/quadratic/fgw_test.py b/tests/solvers/quadratic/fgw_test.py
index c80e64006..380f68fbf 100644
--- a/tests/solvers/quadratic/fgw_test.py
+++ b/tests/solvers/quadratic/fgw_test.py
@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for the Fused Gromov Wasserstein."""
 from typing import Literal, Tuple, Union
 
 import jax
diff --git a/tests/solvers/quadratic/gw_barycenter_test.py b/tests/solvers/quadratic/gw_barycenter_test.py
index 889a3fd13..d16892c3c 100644
--- a/tests/solvers/quadratic/gw_barycenter_test.py
+++ b/tests/solvers/quadratic/gw_barycenter_test.py
@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for Gromov-Wasserstein barycenter."""
 from typing import Any, Optional, Sequence, Tuple
 
 import jax
diff --git a/tests/solvers/quadratic/gw_test.py b/tests/solvers/quadratic/gw_test.py
index 819b9cffb..b7d97a32c 100644
--- a/tests/solvers/quadratic/gw_test.py
+++ b/tests/solvers/quadratic/gw_test.py
@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for the Gromov Wasserstein."""
 from typing import Tuple, Union
 
 import jax
diff --git a/tests/tools/gaussian_mixture/fit_gmm_pair_test.py b/tests/tools/gaussian_mixture/fit_gmm_pair_test.py
index 819e98f77..22b1d430e 100644
--- a/tests/tools/gaussian_mixture/fit_gmm_pair_test.py
+++ b/tests/tools/gaussian_mixture/fit_gmm_pair_test.py
@@ -11,8 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for fit_gmm_pair."""
-
 import jax
 import jax.numpy as jnp
 import pytest
diff --git a/tests/tools/gaussian_mixture/fit_gmm_test.py b/tests/tools/gaussian_mixture/fit_gmm_test.py
index 811c1158e..4ef662d5d 100644
--- a/tests/tools/gaussian_mixture/fit_gmm_test.py
+++ b/tests/tools/gaussian_mixture/fit_gmm_test.py
@@ -11,8 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for fit_gmm_pair."""
-
 import jax
 import jax.numpy as jnp
 import jax.test_util
diff --git a/tests/tools/gaussian_mixture/gaussian_mixture_pair_test.py b/tests/tools/gaussian_mixture/gaussian_mixture_pair_test.py
index d8b3d4f20..3865cbd86 100644
--- a/tests/tools/gaussian_mixture/gaussian_mixture_pair_test.py
+++ b/tests/tools/gaussian_mixture/gaussian_mixture_pair_test.py
@@ -11,8 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for gaussian_mixture_pair."""
-
 import jax
 import jax.numpy as jnp
 import numpy as np
diff --git a/tests/tools/gaussian_mixture/gaussian_mixture_test.py b/tests/tools/gaussian_mixture/gaussian_mixture_test.py
index 92d6650a4..4a23b4b1c 100644
--- a/tests/tools/gaussian_mixture/gaussian_mixture_test.py
+++ b/tests/tools/gaussian_mixture/gaussian_mixture_test.py
@@ -11,8 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for gaussian_mixture."""
-
 import jax
 import jax.numpy as jnp
 import numpy as np
diff --git a/tests/tools/gaussian_mixture/gaussian_test.py b/tests/tools/gaussian_mixture/gaussian_test.py
index 9f32420d7..2c412cc32 100644
--- a/tests/tools/gaussian_mixture/gaussian_test.py
+++ b/tests/tools/gaussian_mixture/gaussian_test.py
@@ -11,8 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for gaussian."""
-
 import jax
 import jax.numpy as jnp
 import numpy as np
diff --git a/tests/tools/gaussian_mixture/linalg_test.py b/tests/tools/gaussian_mixture/linalg_test.py
index 2ae719524..cf9af762e 100644
--- a/tests/tools/gaussian_mixture/linalg_test.py
+++ b/tests/tools/gaussian_mixture/linalg_test.py
@@ -11,8 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for linalg."""
-
 import jax
 import jax.numpy as jnp
 import numpy as np
diff --git a/tests/tools/gaussian_mixture/probabilities_test.py b/tests/tools/gaussian_mixture/probabilities_test.py
index 601333b12..d6db75008 100644
--- a/tests/tools/gaussian_mixture/probabilities_test.py
+++ b/tests/tools/gaussian_mixture/probabilities_test.py
@@ -11,8 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for probabilities."""
-
 import jax
 import jax.numpy as jnp
 import numpy as np
diff --git a/tests/tools/gaussian_mixture/scale_tril_test.py b/tests/tools/gaussian_mixture/scale_tril_test.py
index 54e02db08..eb42cdbcb 100644
--- a/tests/tools/gaussian_mixture/scale_tril_test.py
+++ b/tests/tools/gaussian_mixture/scale_tril_test.py
@@ -11,8 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for ScaleTriL."""
-
 import jax
 import jax.numpy as jnp
 import numpy as np
diff --git a/tests/tools/k_means_test.py b/tests/tools/k_means_test.py
index 5605f1bfc..91c56839f 100644
--- a/tests/tools/k_means_test.py
+++ b/tests/tools/k_means_test.py
@@ -1,3 +1,16 @@
+# 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
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 import os
 import sys
 from typing import Any, Literal, Optional, Tuple, Union
diff --git a/tests/tools/segment_sinkhorn_test.py b/tests/tools/segment_sinkhorn_test.py
index 37c4a8fa7..111af8aae 100644
--- a/tests/tools/segment_sinkhorn_test.py
+++ b/tests/tools/segment_sinkhorn_test.py
@@ -11,8 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for Segmented Sinkhorn."""
-
 import jax
 import jax.numpy as jnp
 import numpy as np
diff --git a/tests/tools/sinkhorn_divergence_test.py b/tests/tools/sinkhorn_divergence_test.py
index 65cdb46c6..0ad7004c8 100644
--- a/tests/tools/sinkhorn_divergence_test.py
+++ b/tests/tools/sinkhorn_divergence_test.py
@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for the Sinkhorn divergence."""
 from typing import Any, Dict, Optional
 
 import jax
diff --git a/tests/tools/soft_sort_test.py b/tests/tools/soft_sort_test.py
index d8242ba91..706014c40 100644
--- a/tests/tools/soft_sort_test.py
+++ b/tests/tools/soft_sort_test.py
@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Tests for the soft sort tools."""
 import functools
 from typing import Tuple
 
diff --git a/tests/utils_test.py b/tests/utils_test.py
index 2e12b7659..768a498b5 100644
--- a/tests/utils_test.py
+++ b/tests/utils_test.py
@@ -1,3 +1,16 @@
+# 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
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 from typing import Optional
 
 import pytest