From 0ff234049b6a630ac91ca2010e27e6b97d897c27 Mon Sep 17 00:00:00 2001
From: Sergei Lebedev <slebedev@google.com>
Date: Mon, 17 Mar 2025 07:49:01 -0700
Subject: [PATCH] Removed trivial docstrings from JAX tests

These docstrings do not make the tests any more clear and typically just duplicate the test module name.

PiperOrigin-RevId: 737611977
---
 jax/experimental/array_serialization/serialization_test.py | 1 -
 jax/experimental/jax2tf/examples/saved_model_main_test.py  | 1 -
 jax/experimental/jax2tf/tests/call_tf_test.py              | 1 -
 jax/experimental/jax2tf/tests/control_flow_ops_test.py     | 1 -
 jax/experimental/jax2tf/tests/shape_poly_test.py           | 1 -
 tests/aot_test.py                                          | 1 -
 tests/api_util_test.py                                     | 1 -
 tests/array_test.py                                        | 1 -
 tests/clear_backends_test.py                               | 1 -
 tests/debug_nans_test.py                                   | 2 --
 tests/garbage_collection_guard_test.py                     | 1 -
 tests/lax_numpy_ufuncs_test.py                             | 2 --
 tests/linalg_test.py                                       | 2 --
 tests/mesh_utils_test.py                                   | 1 -
 tests/mosaic/gpu_test.py                                   | 1 -
 tests/mosaic/profiler_cupti_test.py                        | 1 -
 tests/nn_test.py                                           | 2 --
 tests/optimizers_test.py                                   | 2 --
 tests/pallas/fuser_block_spec_test.py                      | 2 --
 tests/pallas/indexing_test.py                              | 2 --
 tests/pallas/ops_test.py                                   | 2 --
 tests/pallas/pallas_error_handling_test.py                 | 1 -
 tests/pallas/tpu_ops_test.py                               | 1 -
 tests/pallas/tpu_pallas_distributed_test.py                | 2 --
 tests/pallas/tpu_pallas_random_test.py                     | 1 -
 tests/pallas/tpu_pallas_state_test.py                      | 1 -
 tests/pallas/tpu_splash_attention_kernel_test.py           | 1 -
 tests/pallas/tpu_splash_attention_mask_test.py             | 1 -
 tests/pickle_test.py                                       | 1 -
 tests/qdwh_test.py                                         | 1 -
 tests/shape_poly_test.py                                   | 1 -
 tests/stack_test.py                                        | 3 ---
 tests/stax_test.py                                         | 2 --
 tests/svd_test.py                                          | 1 -
 tests/transfer_guard_test.py                               | 1 -
 35 files changed, 47 deletions(-)

diff --git a/jax/experimental/array_serialization/serialization_test.py b/jax/experimental/array_serialization/serialization_test.py
index 6ec621d68..9f4539fc6 100644
--- a/jax/experimental/array_serialization/serialization_test.py
+++ b/jax/experimental/array_serialization/serialization_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 serialization and deserialization of GDA."""
 
 import asyncio
 import math
diff --git a/jax/experimental/jax2tf/examples/saved_model_main_test.py b/jax/experimental/jax2tf/examples/saved_model_main_test.py
index aa6be0cff..5d6982179 100644
--- a/jax/experimental/jax2tf/examples/saved_model_main_test.py
+++ b/jax/experimental/jax2tf/examples/saved_model_main_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 mnist_lib, saved_model_lib, saved_model_main."""
 
 import os
 from absl import flags
diff --git a/jax/experimental/jax2tf/tests/call_tf_test.py b/jax/experimental/jax2tf/tests/call_tf_test.py
index 0cde96aeb..4647a16d7 100644
--- a/jax/experimental/jax2tf/tests/call_tf_test.py
+++ b/jax/experimental/jax2tf/tests/call_tf_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 call_tf."""
 
 from collections.abc import Callable
 import contextlib
diff --git a/jax/experimental/jax2tf/tests/control_flow_ops_test.py b/jax/experimental/jax2tf/tests/control_flow_ops_test.py
index c66a6d696..3b39c8752 100644
--- a/jax/experimental/jax2tf/tests/control_flow_ops_test.py
+++ b/jax/experimental/jax2tf/tests/control_flow_ops_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 jax2tf conversion for control-flow primitives."""
 
 from absl.testing import absltest
 
diff --git a/jax/experimental/jax2tf/tests/shape_poly_test.py b/jax/experimental/jax2tf/tests/shape_poly_test.py
index 3f8320fd8..09da97e84 100644
--- a/jax/experimental/jax2tf/tests/shape_poly_test.py
+++ b/jax/experimental/jax2tf/tests/shape_poly_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 shape-polymorphic jax2tf conversion."""
 
 from __future__ import annotations
 
diff --git a/tests/aot_test.py b/tests/aot_test.py
index 1245967f8..daaeb8417 100644
--- a/tests/aot_test.py
+++ b/tests/aot_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 AOT compilation."""
 
 import contextlib
 import unittest
diff --git a/tests/api_util_test.py b/tests/api_util_test.py
index e34611c6e..26cca4e74 100644
--- a/tests/api_util_test.py
+++ b/tests/api_util_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 jax.api_util."""
 
 import itertools as it
 from absl.testing import absltest
diff --git a/tests/array_test.py b/tests/array_test.py
index 4184c835a..cc8990828 100644
--- a/tests/array_test.py
+++ b/tests/array_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 Array."""
 
 import contextlib
 import math
diff --git a/tests/clear_backends_test.py b/tests/clear_backends_test.py
index 9ea9cac3a..6e98e7293 100644
--- a/tests/clear_backends_test.py
+++ b/tests/clear_backends_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 release_backend_clients."""
 
 from absl.testing import absltest
 import jax
diff --git a/tests/debug_nans_test.py b/tests/debug_nans_test.py
index d3dcfb2e7..c80d23c41 100644
--- a/tests/debug_nans_test.py
+++ b/tests/debug_nans_test.py
@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for --debug_nans."""
-
 from absl.testing import absltest
 
 import jax
diff --git a/tests/garbage_collection_guard_test.py b/tests/garbage_collection_guard_test.py
index 5c34c6de2..f833c9e00 100644
--- a/tests/garbage_collection_guard_test.py
+++ b/tests/garbage_collection_guard_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 garbage allocation guard."""
 
 import gc
 import weakref
diff --git a/tests/lax_numpy_ufuncs_test.py b/tests/lax_numpy_ufuncs_test.py
index 20a1a58a9..fd5050a58 100644
--- a/tests/lax_numpy_ufuncs_test.py
+++ b/tests/lax_numpy_ufuncs_test.py
@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for jax.numpy.ufunc and its methods."""
-
 import itertools
 from functools import partial
 
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
index 2355449fc..feab105cc 100644
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for the LAPAX linear algebra module."""
-
 from functools import partial
 import itertools
 from typing import Iterator
diff --git a/tests/mesh_utils_test.py b/tests/mesh_utils_test.py
index 4f1b1fb03..136b50794 100644
--- a/tests/mesh_utils_test.py
+++ b/tests/mesh_utils_test.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Tests for mesh utils."""
 
 import collections
 from collections.abc import Sequence
diff --git a/tests/mosaic/gpu_test.py b/tests/mosaic/gpu_test.py
index 6426f8006..1f43b46dc 100644
--- a/tests/mosaic/gpu_test.py
+++ b/tests/mosaic/gpu_test.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Tests for Mosaic GPU DSL functions and utilities."""
 
 from collections.abc import Sequence
 import dataclasses
diff --git a/tests/mosaic/profiler_cupti_test.py b/tests/mosaic/profiler_cupti_test.py
index c6f3d23b4..f3dcf71ca 100644
--- a/tests/mosaic/profiler_cupti_test.py
+++ b/tests/mosaic/profiler_cupti_test.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-"""Tests for Mosaic GPU CUPTI-based profiler."""
 
 from absl.testing import absltest, parameterized
 import jax
diff --git a/tests/nn_test.py b/tests/nn_test.py
index 9df315967..ed016ec34 100644
--- a/tests/nn_test.py
+++ b/tests/nn_test.py
@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for nn module."""
-
 import collections
 from functools import partial
 import itertools
diff --git a/tests/optimizers_test.py b/tests/optimizers_test.py
index c4eca0707..2e027e615 100644
--- a/tests/optimizers_test.py
+++ b/tests/optimizers_test.py
@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for the optimizers module."""
-
 import functools
 
 from absl.testing import absltest
diff --git a/tests/pallas/fuser_block_spec_test.py b/tests/pallas/fuser_block_spec_test.py
index d87bfd30e..1b3a21587 100644
--- a/tests/pallas/fuser_block_spec_test.py
+++ b/tests/pallas/fuser_block_spec_test.py
@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for pull block spec."""
-
 from absl.testing import absltest
 from absl.testing import parameterized
 import jax
diff --git a/tests/pallas/indexing_test.py b/tests/pallas/indexing_test.py
index 13f9880e6..c3f3fa6e8 100644
--- a/tests/pallas/indexing_test.py
+++ b/tests/pallas/indexing_test.py
@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for Pallas indexing logic and abstractions."""
-
 from __future__ import annotations
 import sys
 import unittest
diff --git a/tests/pallas/ops_test.py b/tests/pallas/ops_test.py
index 27ec2dc48..0fc375bf6 100644
--- a/tests/pallas/ops_test.py
+++ b/tests/pallas/ops_test.py
@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for common JAX operations within pallas_call."""
-
 from collections.abc import Sequence
 import functools
 import itertools
diff --git a/tests/pallas/pallas_error_handling_test.py b/tests/pallas/pallas_error_handling_test.py
index 34b0ff149..cd5ceecfc 100644
--- a/tests/pallas/pallas_error_handling_test.py
+++ b/tests/pallas/pallas_error_handling_test.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for Pallas error handling."""
 import functools
 import traceback
 
diff --git a/tests/pallas/tpu_ops_test.py b/tests/pallas/tpu_ops_test.py
index c4d600d23..c8def2627 100644
--- a/tests/pallas/tpu_ops_test.py
+++ b/tests/pallas/tpu_ops_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 TPU specific operations within pallas_call."""
 
 import functools
 import math
diff --git a/tests/pallas/tpu_pallas_distributed_test.py b/tests/pallas/tpu_pallas_distributed_test.py
index a329478a1..f7d7daf18 100644
--- a/tests/pallas/tpu_pallas_distributed_test.py
+++ b/tests/pallas/tpu_pallas_distributed_test.py
@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for distributed pallas TPU operations."""
-
 import functools
 import os
 import tempfile
diff --git a/tests/pallas/tpu_pallas_random_test.py b/tests/pallas/tpu_pallas_random_test.py
index c6027cd04..ca8edf7a2 100644
--- a/tests/pallas/tpu_pallas_random_test.py
+++ b/tests/pallas/tpu_pallas_random_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 random ops in Pallas + Mosaic."""
 
 from absl.testing import absltest
 from absl.testing import parameterized
diff --git a/tests/pallas/tpu_pallas_state_test.py b/tests/pallas/tpu_pallas_state_test.py
index ab3a82dab..46f98c087 100644
--- a/tests/pallas/tpu_pallas_state_test.py
+++ b/tests/pallas/tpu_pallas_state_test.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for Pallas mesh API."""
 import functools
 from absl.testing import absltest
 import jax
diff --git a/tests/pallas/tpu_splash_attention_kernel_test.py b/tests/pallas/tpu_splash_attention_kernel_test.py
index c5943fa27..dfe0bcc0d 100644
--- a/tests/pallas/tpu_splash_attention_kernel_test.py
+++ b/tests/pallas/tpu_splash_attention_kernel_test.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for splash_attention."""
 from __future__ import annotations
 
 from collections.abc import Callable
diff --git a/tests/pallas/tpu_splash_attention_mask_test.py b/tests/pallas/tpu_splash_attention_mask_test.py
index f23baf2a0..f39b4d839 100644
--- a/tests/pallas/tpu_splash_attention_mask_test.py
+++ b/tests/pallas/tpu_splash_attention_mask_test.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for splash_attention_masks."""
 from __future__ import annotations
 
 from absl.testing import absltest
diff --git a/tests/pickle_test.py b/tests/pickle_test.py
index d2bc89f44..185eebd90 100644
--- a/tests/pickle_test.py
+++ b/tests/pickle_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 interoperability between JAX and pickling libraries."""
 
 import pickle
 import unittest
diff --git a/tests/qdwh_test.py b/tests/qdwh_test.py
index 705c96f00..91cc3a51f 100644
--- a/tests/qdwh_test.py
+++ b/tests/qdwh_test.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License
 
-"""Tests for the library of QDWH-based polar decomposition."""
 import functools
 
 from absl.testing import absltest
diff --git a/tests/shape_poly_test.py b/tests/shape_poly_test.py
index 67f8e49b3..6d1ffe744 100644
--- a/tests/shape_poly_test.py
+++ b/tests/shape_poly_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 shape-polymorphic export."""
 
 from __future__ import annotations
 
diff --git a/tests/stack_test.py b/tests/stack_test.py
index 655a42571..aa1a02793 100644
--- a/tests/stack_test.py
+++ b/tests/stack_test.py
@@ -12,9 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
-"""Tests for stack."""
-
 from absl.testing import absltest
 
 import jax
diff --git a/tests/stax_test.py b/tests/stax_test.py
index e21300ddd..8c38820d2 100644
--- a/tests/stax_test.py
+++ b/tests/stax_test.py
@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Tests for Stax library."""
-
 from absl.testing import absltest
 
 import numpy as np
diff --git a/tests/svd_test.py b/tests/svd_test.py
index b349c3ca6..97f8176f8 100644
--- a/tests/svd_test.py
+++ b/tests/svd_test.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License
 
-"""Tests for the library of QDWH-based singular value decomposition."""
 import functools
 
 import jax
diff --git a/tests/transfer_guard_test.py b/tests/transfer_guard_test.py
index 6a255b0a1..740ef55ff 100644
--- a/tests/transfer_guard_test.py
+++ b/tests/transfer_guard_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 transfer guards."""
 
 import contextlib
 import pickle