diff --git a/dm_pix/_src/api_test.py b/dm_pix/_src/api_test.py
index 2f2f973..7b0aaef 100644
--- a/dm_pix/_src/api_test.py
+++ b/dm_pix/_src/api_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 dm_pix API."""
-
 import inspect
 
 from absl.testing import absltest
diff --git a/dm_pix/_src/augment_test.py b/dm_pix/_src/augment_test.py
index 7554157..df6f3d9 100644
--- a/dm_pix/_src/augment_test.py
+++ b/dm_pix/_src/augment_test.py
@@ -11,9 +11,8 @@
 # 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 dm_pix._src.augment."""
-
 import functools
+import os
 
 from absl.testing import absltest
 from absl.testing import parameterized
@@ -608,5 +607,6 @@ def test_pad_to_size_batch_chw_shape(self):
 
 
 if __name__ == "__main__":
+  os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
   jax.config.update("jax_default_matmul_precision", "float32")
   absltest.main()
diff --git a/dm_pix/_src/color_conversion_test.py b/dm_pix/_src/color_conversion_test.py
index 4c97583..1d53edd 100644
--- a/dm_pix/_src/color_conversion_test.py
+++ b/dm_pix/_src/color_conversion_test.py
@@ -11,11 +11,10 @@
 # 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 dm_pix._src.color_conversion."""
-
 import colorsys
 import enum
 import functools
+import os
 from typing import Sequence
 
 from absl.testing import parameterized
@@ -285,4 +284,5 @@ def test_grayscale(self, test_images, keep_dims, channel_last):
 
 
 if __name__ == "__main__":
+  os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
   tf.test.main()
diff --git a/dm_pix/_src/depth_and_space_test.py b/dm_pix/_src/depth_and_space_test.py
index d69a87b..a4116b5 100644
--- a/dm_pix/_src/depth_and_space_test.py
+++ b/dm_pix/_src/depth_and_space_test.py
@@ -11,7 +11,7 @@
 # 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 dm_pix._src.depth_and_space."""
+import os
 
 from absl.testing import parameterized
 import chex
@@ -46,4 +46,5 @@ def test_space_to_depth(self, input_shape, block_size):
 
 
 if __name__ == "__main__":
+  os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
   tf.test.main()
diff --git a/dm_pix/_src/interpolation_test.py b/dm_pix/_src/interpolation_test.py
index 7416fb8..4343cda 100644
--- a/dm_pix/_src/interpolation_test.py
+++ b/dm_pix/_src/interpolation_test.py
@@ -11,9 +11,8 @@
 # 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 dm_pix._src.interpolation."""
-
 import itertools
+import os
 from typing import Sequence, Tuple
 
 from absl.testing import absltest
@@ -28,7 +27,8 @@
 
 
 def _prepare_inputs(
-    shape_output_coordinates: Tuple[int]) -> Tuple[jnp.ndarray, jnp.ndarray]:
+    shape_output_coordinates: Tuple[int, ...],
+) -> Tuple[jnp.ndarray, jnp.ndarray]:
   """Returns the volume and coordinates to be used in the function under test.
 
   Args:
@@ -189,4 +189,5 @@ def test_flat_nd_linear_interpolate_constant(self, shape_coordinates, cval):
 
 
 if __name__ == "__main__":
+  os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
   absltest.main()
diff --git a/dm_pix/_src/metrics_test.py b/dm_pix/_src/metrics_test.py
index 46b9cd4..c3e6b09 100644
--- a/dm_pix/_src/metrics_test.py
+++ b/dm_pix/_src/metrics_test.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import functools
+import os
 
 from absl.testing import absltest
 import chex
@@ -206,4 +207,5 @@ def test_ssim_ignore_nans(self):
 
 
 if __name__ == "__main__":
+  os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"
   absltest.main()
diff --git a/dm_pix/_src/patch_test.py b/dm_pix/_src/patch_test.py
index 6c39b34..0ad9bd7 100644
--- a/dm_pix/_src/patch_test.py
+++ b/dm_pix/_src/patch_test.py
@@ -11,9 +11,8 @@
 # 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 dm_pix._src.patch."""
-
 import functools
+import os
 
 from absl.testing import absltest
 from absl.testing import parameterized
@@ -202,4 +201,5 @@ def test_extract_patches_raises(self, sizes, strides, rates, padding):
 
 
 if __name__ == '__main__':
+  os.environ['XLA_PYTHON_CLIENT_PREALLOCATE'] = 'false'
   absltest.main()