Add tensorflow and pytorch random seed setting.

Tests and README are included.

Add support for setting tensorflow and pytorch random seed. It can help detect flaky tests due to the randomness, provide reproducibility.

Co-Authored-By: Saikat Dutta <[email protected]>

Author: 2 people

README.rst

@@ -67,6 +67,10 @@ All of these features are on by default but can be disabled with flags.
   .. |numpy.random| replace:: ``numpy.random``
   __ https://numpy.org/doc/stable/reference/random/index.html
 
+* If `TensorFlow <https://www.tensorflow.org/>`_ is installed, its random seed in ``tensorflow.random`` is reset at the start of every test.
+
+* If `PyTorch <https://pytorch.org/>`_ is installed, its random seed is reset at the start of every test. The random seed of each test is recorded, and can play a role in detecting flaky tests.
+
 * If additional random generators are used, they can be registered under the
   ``pytest_randomly.random_seeder``
   `entry point <https://packaging.python.org/specifications/entry-points/>`_ and

src/pytest_randomly/__init__.py

@@ -67,6 +67,22 @@
 except ImportError:  # pragma: no cover
     have_numpy = False
 
+# tensorflow
+try:
+    import tensorflow as tf
+
+    have_tensorflow = True
+except ImportError:  # pragma: no cover
+    have_tensorflow = False
+
+# pytorch
+try:
+    import torch
+
+    have_pytorch = True
+except ImportError:  # pragma: no cover
+    have_pytorch = False
+
 
 default_seed = random.Random().getrandbits(32)
 
@@ -180,6 +196,17 @@ def _reseed(config: Config, offset: int = 0) -> int:
         else:
             np_random.set_state(np_random_states[numpy_seed])
 
+    if have_tensorflow:  # pragma: no branch
+        tf.random.set_seed(seed)
+        # TensorFlow 1.x compatibility
+        if hasattr(tf, 'compat'):
+            tf.compat.v1.set_random_seed(seed)
+
+    if have_pytorch:  # pragma: no branch
+        torch.manual_seed(seed)
+        if torch.cuda.is_available():  # Also seed CUDA if available
+            torch.cuda.manual_seed_all(seed)
+
     if entrypoint_reseeds is None:
         eps = entry_points(group="pytest_randomly.random_seeder")
         entrypoint_reseeds = [e.load() for e in eps]

tests/test_pytest_randomly.py

@@ -727,6 +727,40 @@ def test_one():
     out.assert_outcomes(passed=1)
 
 
+def test_tensorflow(ourtester):
+    ourtester.makepyfile(
+        test_one="""
+        import tensorflow as tf
+
+        def test_one():
+            assert tf.random.uniform([]) == tf.constant(0.16513085, dtype=tf.float32)
+
+        def test_two():
+            assert tf.random.uniform([]) == tf.constant(0.16513085, dtype=tf.float32)
+        """
+    )
+
+    out = ourtester.runpytest("--randomly-seed=1")
+    out.assert_outcomes(passed=2)
+
+
+def test_pytorch(ourtester):
+    ourtester.makepyfile(
+        test_one="""
+        import torch
+
+        def test_one():
+            assert torch.rand(1) == torch.tensor([0.757631599903106689453125])
+
+        def test_two():
+            assert torch.rand(1) == torch.tensor([0.757631599903106689453125])
+        """
+    )
+
+    out = ourtester.runpytest("--randomly-seed=1")
+    out.assert_outcomes(passed=2)
+
+
 def test_failing_import(testdir):
     """Test with pytest raising CollectError or ImportError.