Dynamically allocate XLA memory for test

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

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()

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()

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()

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()

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()

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()